diff --git a/be/src/core/data_type_serde/data_type_variant_v2_serde.cpp b/be/src/core/data_type_serde/data_type_variant_v2_serde.cpp index e5926b2470aef1..f8a274851de4b9 100644 --- a/be/src/core/data_type_serde/data_type_variant_v2_serde.cpp +++ b/be/src/core/data_type_serde/data_type_variant_v2_serde.cpp @@ -18,6 +18,7 @@ #include "core/data_type_serde/data_type_variant_v2_serde.h" #include +#include #include #include @@ -182,6 +183,133 @@ void preflight_json(const IColumn& column, size_t start, size_t end, }); } +void validate_binary_variant_primitive(VariantPrimitiveId primitive_id) { + switch (primitive_id) { + case VariantPrimitiveId::NULL_VALUE: + case VariantPrimitiveId::TRUE_VALUE: + case VariantPrimitiveId::FALSE_VALUE: + case VariantPrimitiveId::INT8: + case VariantPrimitiveId::INT16: + case VariantPrimitiveId::INT32: + case VariantPrimitiveId::INT64: + case VariantPrimitiveId::DOUBLE: + case VariantPrimitiveId::DECIMAL4: + case VariantPrimitiveId::DECIMAL8: + case VariantPrimitiveId::DECIMAL16: + case VariantPrimitiveId::DATE: + case VariantPrimitiveId::TIMESTAMP_MICROS: + case VariantPrimitiveId::TIMESTAMP_NTZ_MICROS: + case VariantPrimitiveId::FLOAT: + case VariantPrimitiveId::BINARY: + case VariantPrimitiveId::STRING: + case VariantPrimitiveId::UUID: + return; + case VariantPrimitiveId::TIME_NTZ_MICROS: + case VariantPrimitiveId::TIMESTAMP_NANOS: + case VariantPrimitiveId::TIMESTAMP_NTZ_NANOS: + throw Exception(ErrorCode::NOT_IMPLEMENTED_ERROR, + "Binary Variant V2 Arrow encoding does not support primitive id {}", + static_cast(primitive_id)); + } + throw Exception(ErrorCode::NOT_IMPLEMENTED_ERROR, + "Binary Variant V2 Arrow encoding does not support unknown primitive id {}", + static_cast(primitive_id)); +} + +void validate_binary_variant_value(VariantRef value, uint32_t depth = 0) { + if (depth > VARIANT_MAX_NESTING_DEPTH) { + throw Exception(ErrorCode::CORRUPTION, "Variant value exceeds maximum nesting depth {}", + VARIANT_MAX_NESTING_DEPTH); + } + const size_t encoded_size = value.value_size(); + if (encoded_size != value.value.size) { + throw Exception(ErrorCode::CORRUPTION, + "Variant value has {} trailing bytes after the encoded value", + value.value.size - encoded_size); + } + + switch (value.basic_type()) { + case VariantBasicType::PRIMITIVE: + validate_binary_variant_primitive(value.primitive_id()); + return; + case VariantBasicType::SHORT_STRING: + return; + case VariantBasicType::OBJECT: + for (uint32_t i = 0; i < value.num_elements(); ++i) { + uint32_t field_id = 0; + VariantRef child = value.object_value_at(i, &field_id); + value.metadata.key_at(field_id); + validate_binary_variant_value(child, depth + 1); + } + return; + case VariantBasicType::ARRAY: + for (uint32_t i = 0; i < value.num_elements(); ++i) { + validate_binary_variant_value(value.array_at(i), depth + 1); + } + return; + } +} + +void require_variant_arrow_status(const arrow::Status& status) { + if (!status.ok()) { + throw Exception(ErrorCode::INTERNAL_ERROR, "Variant V2 Arrow append failed: {}", + status.ToString()); + } +} + +Status write_binary_variant_arrow(const IColumn& column, const NullMap* null_map, + arrow::StructBuilder& builder, size_t start, size_t end) { + // StructBuilder::type() returns a shared_ptr by value. Keep that owner alive while using the + // cast reference; otherwise the reference would dangle as soon as the temporary is destroyed. + const auto builder_type = builder.type(); + const auto& struct_type = assert_cast(*builder_type); + if (struct_type.num_fields() != 2 || struct_type.field(0)->name() != "value" || + struct_type.field(1)->name() != "metadata" || + struct_type.field(0)->type()->id() != arrow::Type::BINARY || + struct_type.field(1)->type()->id() != arrow::Type::BINARY) { + return Status::InvalidArgument( + "Binary Variant V2 Arrow type must be " + "struct, got {}", + struct_type.ToString()); + } + auto* value_builder = dynamic_cast(builder.field_builder(0)); + auto* metadata_builder = dynamic_cast(builder.field_builder(1)); + if (value_builder == nullptr || metadata_builder == nullptr) { + return Status::InvalidArgument("Binary Variant V2 Arrow child builders must be binary"); + } + + // Consumers of the binary Variant V2 Arrow representation copy these two buffers without + // inspecting them. Validate once at the serialization boundary so a writer cannot commit + // malformed or unsupported bytes. + const auto outer_nulls = forced_nulls(null_map); + visit_variant_v2_values( + column, start, end, outer_nulls, + [&](size_t) { require_variant_arrow_status(builder.AppendNull()); }, + [&](size_t row, VariantRef value) { + try { + constexpr size_t BINARY_VARIANT_SIZE_LIMIT = 128 * 1024 * 1024; + if (value.value.size > BINARY_VARIANT_SIZE_LIMIT || + value.metadata.size > BINARY_VARIANT_SIZE_LIMIT) { + throw Exception(ErrorCode::INVALID_ARGUMENT, + "exceeds the 128 MiB value/metadata limit"); + } + value.metadata.validate(); + validate_binary_variant_value(value); + } catch (const Exception& e) { + throw Exception(e.code(), "Binary Variant V2 row {} is incompatible: {}", row, + e.what()); + } + require_variant_arrow_status(builder.Append()); + require_variant_arrow_status( + value_builder->Append(reinterpret_cast(value.value.data), + cast_set(value.value.size))); + require_variant_arrow_status(metadata_builder->Append( + reinterpret_cast(value.metadata.data), + cast_set(value.metadata.size))); + }); + return Status::OK(); +} + } // namespace DataTypeVariantV2SerDe::DataTypeVariantV2SerDe(int nesting_level) : DataTypeSerDe(nesting_level) {} @@ -553,6 +681,11 @@ Status DataTypeVariantV2SerDe::write_column_to_arrow(const IColumn& column, cons assert_cast(*array_builder), first, last, options); } + if (array_builder->type()->id() == arrow::Type::STRUCT) { + return write_binary_variant_arrow(column, null_map, + assert_cast(*array_builder), + first, last); + } return Status::InvalidArgument("Unsupported arrow type for variant column: {}", array_builder->type()->name()); }); diff --git a/be/src/exec/operator/exchange_sink_operator.cpp b/be/src/exec/operator/exchange_sink_operator.cpp index 449c71e3339281..f7f8ce45d0ff9b 100644 --- a/be/src/exec/operator/exchange_sink_operator.cpp +++ b/be/src/exec/operator/exchange_sink_operator.cpp @@ -28,6 +28,7 @@ #include #include +#include "agent/be_exec_version_manager.h" #include "common/status.h" #include "core/column/column_const.h" #include "exec/exchange/exchange_writer.h" @@ -35,6 +36,7 @@ #include "exec/operator/exchange_sink_buffer.h" #include "exec/operator/operator.h" #include "exec/operator/sort_source_operator.h" +#include "exec/partitioner/external/external_table_sink_hash_partitioner.h" #include "exec/pipeline/dependency.h" #include "exec/pipeline/pipeline_fragment_context.h" #include "exec/sink/scale_writer_partitioning_exchanger.hpp" @@ -98,7 +100,8 @@ Status ExchangeSinkLocalState::init(RuntimeState* state, LocalSinkStateInfo& inf _part_type = p._part_type; // Shuffle the channels randomly if (_part_type == TPartitionType::UNPARTITIONED || _part_type == TPartitionType::RANDOM || - _part_type == TPartitionType::HIVE_TABLE_SINK_UNPARTITIONED) { + _part_type == TPartitionType::HIVE_TABLE_SINK_UNPARTITIONED || + _part_type == TPartitionType::EXTERNAL_TABLE_SINK_UNPARTITIONED) { std::random_device rd; std::mt19937 g(rd()); shuffle(channels.begin(), channels.end(), g); @@ -175,6 +178,27 @@ Status ExchangeSinkLocalState::init(RuntimeState* state, LocalSinkStateInfo& inf RETURN_IF_ERROR(_partitioner->prepare(state, p._row_desc)); custom_profile()->add_info_string( "Partitioner", fmt::format("ScaleWriterPartitioner({})", _partition_count)); + } else if (_part_type == TPartitionType::EXTERNAL_TABLE_SINK_HASH_PARTITIONED) { + if (state->be_exec_version() < SUPPORT_EXTERNAL_TABLE_SINK_HASH_VERSION) { + return Status::NotSupported( + "External table sink hash exchange requires BE execution version {}, actual {}", + SUPPORT_EXTERNAL_TABLE_SINK_HASH_VERSION, state->be_exec_version()); + } + if (!p._has_external_table_sink_hash_partition_info) { + return Status::InternalError("External table sink hash partition info is missing"); + } + _partition_count = channels.size(); + const bool use_crc32c = _state->query_options().__isset.enable_new_shuffle_hash_method && + _state->query_options().enable_new_shuffle_hash_method; + const ShuffleHashMethod hash_method = + use_crc32c ? ShuffleHashMethod::CRC32C : ShuffleHashMethod::CRC32; + _partitioner = std::make_unique( + _partition_count, hash_method, p._external_table_sink_hash_partition_info); + RETURN_IF_ERROR(_partitioner->init(p._texprs)); + RETURN_IF_ERROR(_partitioner->prepare(state, p._row_desc)); + custom_profile()->add_info_string( + "Partitioner", + fmt::format("ExternalTableSinkHashPartitioner({})", _partition_count)); } else if (_part_type == TPartitionType::MERGE_PARTITIONED) { if (!p._has_merge_partition_info) { return Status::InternalError("Merge partition info is missing"); @@ -272,6 +296,7 @@ Status ExchangeSinkLocalState::open(RuntimeState* state) { if (_part_type == TPartitionType::HASH_PARTITIONED || _part_type == TPartitionType::BUCKET_SHFFULE_HASH_PARTITIONED || _part_type == TPartitionType::HIVE_TABLE_SINK_HASH_PARTITIONED || + _part_type == TPartitionType::EXTERNAL_TABLE_SINK_HASH_PARTITIONED || _part_type == TPartitionType::OLAP_TABLE_SINK_HASH_PARTITIONED || _part_type == TPartitionType::MERGE_PARTITIONED) { RETURN_IF_ERROR(_partitioner->open(state)); @@ -322,6 +347,8 @@ ExchangeSinkOperatorX::ExchangeSinkOperatorX( sink.output_partition.type == TPartitionType::BUCKET_SHFFULE_HASH_PARTITIONED || sink.output_partition.type == TPartitionType::HIVE_TABLE_SINK_HASH_PARTITIONED || sink.output_partition.type == TPartitionType::HIVE_TABLE_SINK_UNPARTITIONED || + sink.output_partition.type == TPartitionType::EXTERNAL_TABLE_SINK_HASH_PARTITIONED || + sink.output_partition.type == TPartitionType::EXTERNAL_TABLE_SINK_UNPARTITIONED || sink.output_partition.type == TPartitionType::MERGE_PARTITIONED); #endif _name = "ExchangeSinkOperatorX"; @@ -333,6 +360,11 @@ ExchangeSinkOperatorX::ExchangeSinkOperatorX( _merge_partition_info = sink.output_partition.merge_partition_info; _has_merge_partition_info = true; } + if (sink.output_partition.__isset.external_table_sink_hash_partition_info) { + _external_table_sink_hash_partition_info = + sink.output_partition.external_table_sink_hash_partition_info; + _has_external_table_sink_hash_partition_info = true; + } if (_part_type != TPartitionType::UNPARTITIONED) { // if the destinations only one dest, we need to use broadcast @@ -535,20 +567,16 @@ Status ExchangeSinkOperatorX::sink_impl(RuntimeState* state, Block* block, bool _part_type == TPartitionType::BUCKET_SHFFULE_HASH_PARTITIONED || _part_type == TPartitionType::OLAP_TABLE_SINK_HASH_PARTITIONED || _part_type == TPartitionType::HIVE_TABLE_SINK_HASH_PARTITIONED || + _part_type == TPartitionType::EXTERNAL_TABLE_SINK_HASH_PARTITIONED || _part_type == TPartitionType::MERGE_PARTITIONED) { RETURN_IF_ERROR(local_state._writer->write(state, block, eos)); - } else if (_part_type == TPartitionType::HIVE_TABLE_SINK_UNPARTITIONED) { + } else if (_part_type == TPartitionType::HIVE_TABLE_SINK_UNPARTITIONED || + _part_type == TPartitionType::EXTERNAL_TABLE_SINK_UNPARTITIONED) { // Control the number of channels according to the flow, thereby controlling the number of table sink writers. RETURN_IF_ERROR(send_to_current_channel()); - _data_processed += block->bytes(); - if (_writer_count < local_state.channels.size()) { - if (_data_processed >= - _writer_count * - config::table_sink_non_partition_write_scaling_data_processed_threshold) { - _writer_count++; - } - } - local_state.current_channel_idx = (local_state.current_channel_idx + 1) % _writer_count; + const auto writer_count = + _update_writer_scaling(block->bytes(), local_state.channels.size()); + local_state.current_channel_idx = (local_state.current_channel_idx + 1) % writer_count; } else { // Range partition // 1. calculate range @@ -587,6 +615,18 @@ Status ExchangeSinkOperatorX::sink_impl(RuntimeState* state, Block* block, bool return final_st; } +size_t ExchangeSinkOperatorX::_update_writer_scaling(size_t block_bytes, size_t max_writer_count) { + LockGuard lock(_writer_scaling_mutex); + _data_processed += block_bytes; + if (_writer_count < max_writer_count && + _data_processed >= + _writer_count * + config::table_sink_non_partition_write_scaling_data_processed_threshold) { + ++_writer_count; + } + return _writer_count; +} + void ExchangeSinkLocalState::register_channels(ExchangeSinkBuffer* buffer) { for (auto& channel : channels) { channel->set_exchange_buffer(buffer); diff --git a/be/src/exec/operator/exchange_sink_operator.h b/be/src/exec/operator/exchange_sink_operator.h index 10351154d1d8cd..417a72da86bf28 100644 --- a/be/src/exec/operator/exchange_sink_operator.h +++ b/be/src/exec/operator/exchange_sink_operator.h @@ -200,6 +200,17 @@ class ExchangeSinkOperatorX MOCK_REMOVE(final) : public DataSinkOperatorX writer_scaling_state_for_test() { + LockGuard lock(_writer_scaling_mutex); + return {_data_processed, _writer_count}; + } +#endif + bool is_serial_operator() const override { return true; } void set_low_memory_mode(RuntimeState* state) override { auto& local_state = get_local_state(state); @@ -239,10 +250,13 @@ class ExchangeSinkOperatorX MOCK_REMOVE(final) : public DataSinkOperatorX _create_buffer( RuntimeState* state, const std::vector& sender_ins_ids); + size_t _update_writer_scaling(size_t block_bytes, size_t max_writer_count); std::shared_ptr _sink_buffer = nullptr; RuntimeState* _state = nullptr; const std::vector _texprs; + TExternalTableSinkHashPartitionInfo _external_table_sink_hash_partition_info; + bool _has_external_table_sink_hash_partition_info = false; TMergePartitionInfo _merge_partition_info; bool _has_merge_partition_info = false; @@ -278,8 +292,9 @@ class ExchangeSinkOperatorX MOCK_REMOVE(final) : public DataSinkOperatorX& partition_exprs, + std::unique_ptr* partition_function) { + if (partition_function == nullptr) { + return Status::InvalidArgument("External partition function output is null"); + } + if (partition_info.partition_function != "direct_hash") { + return Status::NotSupported("Unsupported external sink partition function '{}'", + partition_info.partition_function); + } + if (partition_info.__isset.partition_function_options && + !partition_info.partition_function_options.empty()) { + return Status::InvalidArgument("Direct hash partition function does not accept options"); + } + auto function = std::make_unique(logical_partition_count, hash_method); + RETURN_IF_ERROR(function->init(partition_exprs)); + *partition_function = std::move(function); + return Status::OK(); +} + +} // namespace doris diff --git a/be/src/exec/partitioner/external/external_partition_function_factory.h b/be/src/exec/partitioner/external/external_partition_function_factory.h new file mode 100644 index 00000000000000..8f7bdee08073e7 --- /dev/null +++ b/be/src/exec/partitioner/external/external_partition_function_factory.h @@ -0,0 +1,34 @@ +// 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. + +#pragma once + +#include + +#include +#include + +#include "exec/partitioner/partitioner.h" + +namespace doris { + +Status create_external_partition_function(const TExternalTableSinkHashPartitionInfo& partition_info, + PartitionerBase::HashValType logical_partition_count, + ShuffleHashMethod hash_method, + const std::vector& partition_exprs, + std::unique_ptr* partition_function); + +} // namespace doris diff --git a/be/src/exec/partitioner/external/external_table_sink_hash_partitioner.cpp b/be/src/exec/partitioner/external/external_table_sink_hash_partitioner.cpp new file mode 100644 index 00000000000000..2d65d481778774 --- /dev/null +++ b/be/src/exec/partitioner/external/external_table_sink_hash_partitioner.cpp @@ -0,0 +1,101 @@ +// 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/partitioner/external/external_table_sink_hash_partitioner.h" + +#include +#include + +#include "common/cast_set.h" +#include "common/config.h" +#include "common/status.h" +#include "exec/partitioner/external/external_partition_function_factory.h" + +namespace doris { + +ExternalTableSinkHashPartitioner::ExternalTableSinkHashPartitioner( + HashValType partition_count, ShuffleHashMethod hash_method, + TExternalTableSinkHashPartitionInfo partition_info) + : PartitionerBase(partition_count), + _hash_method(hash_method), + _partition_info(std::move(partition_info)), + _logical_partition_count(partition_count) {} + +Status ExternalTableSinkHashPartitioner::init(const std::vector& texprs) { + if (_partition_info.writer_assignment == TExternalTableSinkWriterAssignment::SKEWED) { + const auto partitions_per_writer = cast_set( + std::max(1, config::table_sink_partition_write_max_partition_nums_per_writer)); + if (_partition_count > std::numeric_limits::max() / partitions_per_writer) { + return Status::InvalidArgument("External sink logical partition count overflows"); + } + _logical_partition_count = _partition_count * partitions_per_writer; + } else if (_partition_info.writer_assignment != TExternalTableSinkWriterAssignment::IDENTITY) { + return Status::InvalidArgument("Unsupported external sink writer assignment {}", + static_cast(_partition_info.writer_assignment)); + } + return create_external_partition_function(_partition_info, _logical_partition_count, + _hash_method, texprs, &_partition_function); +} + +Status ExternalTableSinkHashPartitioner::prepare(RuntimeState* state, + const RowDescriptor& row_desc) { + return _partition_function->prepare(state, row_desc); +} + +Status ExternalTableSinkHashPartitioner::open(RuntimeState* state) { + RETURN_IF_ERROR(_partition_function->open(state)); + if (_partition_info.writer_assignment == TExternalTableSinkWriterAssignment::IDENTITY) { + _writer_assigner = std::make_unique(_partition_count); + } else { + _writer_assigner = std::make_unique( + cast_set(_logical_partition_count), cast_set(_partition_count), 1, + scale_writer_threshold_by_task( + config::table_sink_partition_write_min_partition_data_processed_rebalance_threshold, + state->task_num()), + scale_writer_threshold_by_task( + config::table_sink_partition_write_min_data_processed_rebalance_threshold, + state->task_num())); + } + return Status::OK(); +} + +Status ExternalTableSinkHashPartitioner::close(RuntimeState* state) { + return _partition_function->close(state); +} + +Status ExternalTableSinkHashPartitioner::do_partitioning(RuntimeState* state, Block* block) const { + RETURN_IF_ERROR(_partition_function->get_partitions(state, block, _logical_partition_count, + _logical_partition_ids)); + return _writer_assigner->assign(_logical_partition_ids, nullptr, block->rows(), block->bytes(), + _channel_ids); +} + +const std::vector& +ExternalTableSinkHashPartitioner::get_channel_ids() const { + return _channel_ids; +} + +Status ExternalTableSinkHashPartitioner::clone(RuntimeState* state, + std::unique_ptr& partitioner) { + auto cloned = std::make_unique(_partition_count, _hash_method, + _partition_info); + RETURN_IF_ERROR(_partition_function->clone(state, cloned->_partition_function)); + cloned->_logical_partition_count = _logical_partition_count; + partitioner = std::move(cloned); + return Status::OK(); +} + +} // namespace doris diff --git a/be/src/exec/partitioner/external/external_table_sink_hash_partitioner.h b/be/src/exec/partitioner/external/external_table_sink_hash_partitioner.h new file mode 100644 index 00000000000000..fe8ce8637af0ba --- /dev/null +++ b/be/src/exec/partitioner/external/external_table_sink_hash_partitioner.h @@ -0,0 +1,51 @@ +// 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. + +#pragma once + +#include + +#include + +#include "exec/partitioner/partitioner.h" +#include "exec/partitioner/writer_assigner.h" + +namespace doris { + +class ExternalTableSinkHashPartitioner final : public PartitionerBase { +public: + ExternalTableSinkHashPartitioner(HashValType partition_count, ShuffleHashMethod hash_method, + TExternalTableSinkHashPartitionInfo partition_info); + + Status init(const std::vector& texprs) override; + Status prepare(RuntimeState* state, const RowDescriptor& row_desc) override; + Status open(RuntimeState* state) override; + Status close(RuntimeState* state) override; + Status do_partitioning(RuntimeState* state, Block* block) const override; + const std::vector& get_channel_ids() const override; + Status clone(RuntimeState* state, std::unique_ptr& partitioner) override; + +private: + ShuffleHashMethod _hash_method; + TExternalTableSinkHashPartitionInfo _partition_info; + HashValType _logical_partition_count; + std::unique_ptr _partition_function; + std::unique_ptr _writer_assigner; + mutable std::vector _logical_partition_ids; + mutable std::vector _channel_ids; +}; + +} // namespace doris diff --git a/be/src/exec/partitioner/partitioner.cpp b/be/src/exec/partitioner/partitioner.cpp index a7290be8c2925b..747f9ed33377cb 100644 --- a/be/src/exec/partitioner/partitioner.cpp +++ b/be/src/exec/partitioner/partitioner.cpp @@ -84,6 +84,51 @@ Status Crc32CHashPartitioner::clone(RuntimeState* state, return _clone_expr_ctxs(state, new_partitioner->_partition_expr_ctxs); } +HashPartitionFunction::HashPartitionFunction(HashValType partition_count, + ShuffleHashMethod hash_method) + : _partition_count(partition_count), _hash_method(hash_method) {} + +Status HashPartitionFunction::init(const std::vector& texprs) { + if (_hash_method == ShuffleHashMethod::CRC32C) { + _partitioner = std::make_unique(_partition_count); + } else { + _partitioner = std::make_unique>(_partition_count); + } + return _partitioner->init(texprs); +} + +Status HashPartitionFunction::prepare(RuntimeState* state, const RowDescriptor& row_desc) { + return _partitioner->prepare(state, row_desc); +} + +Status HashPartitionFunction::open(RuntimeState* state) { + return _partitioner->open(state); +} + +Status HashPartitionFunction::close(RuntimeState* state) { + return _partitioner->close(state); +} + +Status HashPartitionFunction::get_partitions(RuntimeState* state, Block* block, + size_t partition_count, + std::vector& partitions) const { + if (partition_count != _partition_count) { + return Status::InvalidArgument("Hash partition count {} does not match planned count {}", + partition_count, _partition_count); + } + RETURN_IF_ERROR(_partitioner->do_partitioning(state, block)); + partitions = _partitioner->get_channel_ids(); + return Status::OK(); +} + +Status HashPartitionFunction::clone(RuntimeState* state, + std::unique_ptr& function) const { + auto cloned = std::make_unique(_partition_count, _hash_method); + RETURN_IF_ERROR(_partitioner->clone(state, cloned->_partitioner)); + function = std::move(cloned); + return Status::OK(); +} + template class Crc32HashPartitioner; template class Crc32HashPartitioner; template class Crc32HashPartitioner; diff --git a/be/src/exec/partitioner/partitioner.h b/be/src/exec/partitioner/partitioner.h index 98607c3623634f..3f6562f8e3bb9f 100644 --- a/be/src/exec/partitioner/partitioner.h +++ b/be/src/exec/partitioner/partitioner.h @@ -55,6 +55,11 @@ class PartitionerBase { const HashValType _partition_count; }; +enum class ShuffleHashMethod { + CRC32, + CRC32C, +}; + class PartitionFunction { public: using HashValType = PartitionerBase::HashValType; @@ -78,9 +83,25 @@ class PartitionFunction { std::unique_ptr& function) const = 0; }; -enum class ShuffleHashMethod { - CRC32, - CRC32C, +// Adapts the standard Doris expression hash partitioner to the composable +// PartitionFunction interface used by sink routing. +class HashPartitionFunction final : public PartitionFunction { +public: + HashPartitionFunction(HashValType partition_count, ShuffleHashMethod hash_method); + + Status init(const std::vector& texprs) override; + Status prepare(RuntimeState* state, const RowDescriptor& row_desc) override; + Status open(RuntimeState* state) override; + Status close(RuntimeState* state) override; + Status get_partitions(RuntimeState* state, Block* block, size_t partition_count, + std::vector& partitions) const override; + HashValType partition_count() const override { return _partition_count; } + Status clone(RuntimeState* state, std::unique_ptr& function) const override; + +private: + HashValType _partition_count; + ShuffleHashMethod _hash_method; + std::unique_ptr _partitioner; }; template diff --git a/be/src/exec/partitioner/writer_assigner.cpp b/be/src/exec/partitioner/writer_assigner.cpp new file mode 100644 index 00000000000000..3245616b1fe36c --- /dev/null +++ b/be/src/exec/partitioner/writer_assigner.cpp @@ -0,0 +1,136 @@ +// 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/partitioner/writer_assigner.h" + +#include "exec/connector/skewed_partition_rebalancer.h" + +namespace doris { + +namespace { +Status validate_assignment_input(const std::vector& partition_ids, + const std::vector* mask, size_t rows) { + if (partition_ids.size() < rows) { + return Status::InvalidArgument("Writer assignment has {} partition ids for {} rows", + partition_ids.size(), rows); + } + if (mask != nullptr && mask->size() < rows) { + return Status::InvalidArgument("Writer assignment mask has {} entries for {} rows", + mask->size(), rows); + } + return Status::OK(); +} +} // namespace + +Status IdentityWriterAssigner::assign(const std::vector& partition_ids, + const std::vector* mask, size_t rows, + size_t /*block_bytes*/, std::vector& writer_ids) { + RETURN_IF_ERROR(validate_assignment_input(partition_ids, mask, rows)); + if (writer_ids.size() != rows && &writer_ids != &partition_ids) { + writer_ids.resize(rows); + } + for (size_t row = 0; row < rows; ++row) { + if (mask != nullptr && (*mask)[row] == 0) { + continue; + } + if (partition_ids[row] >= _writer_count) { + return Status::InvalidArgument("Logical partition {} exceeds writer count {}", + partition_ids[row], _writer_count); + } + writer_ids[row] = partition_ids[row]; + } + return Status::OK(); +} + +SkewedWriterAssigner::SkewedWriterAssigner(int partition_count, int task_count, + int task_bucket_count, + long min_partition_data_processed_rebalance_threshold, + long min_data_processed_rebalance_threshold) + : _rebalancer(std::make_unique( + partition_count, task_count, task_bucket_count, + min_partition_data_processed_rebalance_threshold, + min_data_processed_rebalance_threshold)), + _writer_count(task_count), + _partition_row_counts(partition_count, 0), + _partition_writer_ids(partition_count, -1), + _partition_writer_indexes(partition_count, 0) {} + +SkewedWriterAssigner::~SkewedWriterAssigner() = default; + +Status SkewedWriterAssigner::assign(const std::vector& partition_ids, + const std::vector* mask, size_t rows, + size_t block_bytes, std::vector& writer_ids) { + RETURN_IF_ERROR(validate_assignment_input(partition_ids, mask, rows)); + if (rows == 0) { + return Status::OK(); + } + if (_partition_row_counts.empty()) { + return Status::InvalidArgument("Skewed writer assignment has no logical partitions"); + } + if (writer_ids.size() != rows && &writer_ids != &partition_ids) { + writer_ids.resize(rows); + } + + std::fill(_partition_row_counts.begin(), _partition_row_counts.end(), 0); + std::fill(_partition_writer_ids.begin(), _partition_writer_ids.end(), -1); + _rebalancer->rebalance(); + + const size_t partition_count = _partition_row_counts.size(); + for (size_t row = 0; row < rows; ++row) { + if (mask != nullptr && (*mask)[row] == 0) { + continue; + } + const uint32_t partition_id = partition_ids[row]; + if (partition_id >= partition_count) { + return Status::InvalidArgument("Logical partition {} exceeds partition count {}", + partition_id, partition_count); + } + _partition_row_counts[partition_id] += 1; + int writer_id = _partition_writer_ids[partition_id]; + if (writer_id == -1) { + writer_id = _get_next_writer_id(partition_id); + if (writer_id < 0 || writer_id >= _writer_count) { + return Status::InternalError("Skewed writer assignment returned invalid writer {}", + writer_id); + } + _partition_writer_ids[partition_id] = writer_id; + } + writer_ids[row] = static_cast(writer_id); + } + + for (size_t partition_id = 0; partition_id < partition_count; ++partition_id) { + if (_partition_row_counts[partition_id] > 0) { + _rebalancer->add_partition_row_count(static_cast(partition_id), + _partition_row_counts[partition_id]); + } + } + _rebalancer->add_data_processed(static_cast(block_bytes)); + return Status::OK(); +} + +int SkewedWriterAssigner::_get_next_writer_id(uint32_t partition_id) { + return _rebalancer->get_task_id(partition_id, _partition_writer_indexes[partition_id]++); +} + +int64_t scale_writer_threshold_by_task(int64_t value, int task_num) { + if (task_num <= 0) { + return value; + } + int64_t scaled = value / task_num; + return scaled == 0 ? value : scaled; +} + +} // namespace doris diff --git a/be/src/exec/partitioner/writer_assigner.h b/be/src/exec/partitioner/writer_assigner.h new file mode 100644 index 00000000000000..7a4237cc3645ea --- /dev/null +++ b/be/src/exec/partitioner/writer_assigner.h @@ -0,0 +1,81 @@ +// 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. + +#pragma once + +#include +#include +#include +#include + +#include "common/status.h" + +namespace doris { +class SkewedPartitionRebalancer; +} + +namespace doris { + +// Maps logical partitions computed by a PartitionFunction to Doris exchange channels. +class WriterAssigner { +public: + virtual ~WriterAssigner() = default; + + virtual Status assign(const std::vector& partition_ids, + const std::vector* mask, size_t rows, size_t block_bytes, + std::vector& writer_ids) = 0; +}; + +// Preserves stable ownership: one logical partition always maps to one writer id. +class IdentityWriterAssigner final : public WriterAssigner { +public: + explicit IdentityWriterAssigner(uint32_t writer_count) : _writer_count(writer_count) {} + + Status assign(const std::vector& partition_ids, const std::vector* mask, + size_t rows, size_t block_bytes, std::vector& writer_ids) override; + +private: + uint32_t _writer_count; +}; + +// Allows a hot logical partition to use multiple writers while retaining the existing +// ScaleWriter affinity and rebalance behavior. +class SkewedWriterAssigner final : public WriterAssigner { +public: + SkewedWriterAssigner(int partition_count, int task_count, int task_bucket_count, + long min_partition_data_processed_rebalance_threshold, + long min_data_processed_rebalance_threshold); + + ~SkewedWriterAssigner() override; + + Status assign(const std::vector& partition_ids, const std::vector* mask, + size_t rows, size_t block_bytes, std::vector& writer_ids) override; + +private: + int _get_next_writer_id(uint32_t partition_id); + + std::unique_ptr _rebalancer; + int _writer_count; + std::vector _partition_row_counts; + std::vector _partition_writer_ids; + std::vector _partition_writer_indexes; +}; + +// Scale table-sink thresholds by local pipeline task count while preserving the historical +// behavior for very small values. +int64_t scale_writer_threshold_by_task(int64_t value, int task_num); + +} // namespace doris diff --git a/be/src/exec/pipeline/pipeline_fragment_context.cpp b/be/src/exec/pipeline/pipeline_fragment_context.cpp index 81091aa2421c26..220c32944bff7d 100644 --- a/be/src/exec/pipeline/pipeline_fragment_context.cpp +++ b/be/src/exec/pipeline/pipeline_fragment_context.cpp @@ -2592,7 +2592,8 @@ void PipelineFragmentContext::_coordinator_callback(const ReportStatusRequest& r PrintThriftNetworkAddress(req.coord_addr), e.what()); } - const bool requires_external_file_ack = params.__isset.iceberg_commit_datas; + const bool requires_external_file_ack = + params.__isset.iceberg_commit_datas || params.__isset.connector_commit_data; if (rpc_status.ok() && requires_external_file_ack && (!res.__isset.external_file_commit_data_accepted || !res.external_file_commit_data_accepted)) { diff --git a/be/src/exec/sink/writer/async_result_writer.cpp b/be/src/exec/sink/writer/async_result_writer.cpp index 46c6cdf159cbf9..6397dc660e0544 100644 --- a/be/src/exec/sink/writer/async_result_writer.cpp +++ b/be/src/exec/sink/writer/async_result_writer.cpp @@ -315,7 +315,7 @@ std::unique_ptr AsyncResultWriter::_get_free_block(doris::Block* block, s template void clear_blocks(moodycamel::ConcurrentQueue& blocks, - RuntimeProfile::Counter* memory_used_counter = nullptr); + RuntimeProfile::Counter* memory_used_counter); void AsyncResultWriter::set_low_memory_mode() { _low_memory_mode = true; clear_blocks(_free_blocks, _memory_used_counter); diff --git a/be/src/exprs/function/cast/variant_v2/cast_array_to_variant.cpp b/be/src/exprs/function/cast/variant_v2/cast_array_to_variant.cpp index b30ad8ef6732cf..26cbdd2b71b9c0 100644 --- a/be/src/exprs/function/cast/variant_v2/cast_array_to_variant.cpp +++ b/be/src/exprs/function/cast/variant_v2/cast_array_to_variant.cpp @@ -137,7 +137,9 @@ Status build_array_node_plan(const ColumnPtr& source, const DataTypePtr& source_ Status build_array_leaf_plan(const ColumnPtr& source, PrimitiveType primitive, ArrayEncodePlan* plan) { - if (primitive == INVALID_TYPE && source->empty()) { + if (primitive == INVALID_TYPE) { + // DataTypeNothing is represented by the element null map, including non-empty + // expressions such as array(NULL). return Status::OK(); } else if (primitive == TYPE_VARIANT) { const auto* variant = check_and_get_column(source.get()); @@ -196,7 +198,7 @@ void append_array_value(const ArrayEncodePlan& plan, size_t index, VariantBatchB } else if (plan.jsonb_leaf != nullptr) { jsonb_to_variant(plan.jsonb_leaf->get_data_at(index), *row); } else { - DORIS_CHECK(false) << "empty Array leaf unexpectedly contains a value"; + DORIS_CHECK(false) << "Array Variant V2 leaf has no encoder"; } return; } diff --git a/be/src/format/transformer/merge_partitioner.cpp b/be/src/format/transformer/merge_partitioner.cpp index 89cf830d6bba53..bc6bc265388853 100644 --- a/be/src/format/transformer/merge_partitioner.cpp +++ b/be/src/format/transformer/merge_partitioner.cpp @@ -33,16 +33,6 @@ namespace doris { -namespace { -int64_t scale_threshold_by_task(int64_t value, int task_num) { - if (task_num <= 0) { - return value; - } - int64_t scaled = value / task_num; - return scaled == 0 ? value : scaled; -} -} // namespace - MergePartitioner::MergePartitioner(size_t partition_count, const TMergePartitionInfo& merge_info, bool use_new_shuffle_hash_method) : PartitionerBase(static_cast(partition_count)), @@ -183,7 +173,7 @@ Status MergePartitioner::do_partitioning(RuntimeState* state, Block* block) cons _insert_writer_count = static_cast(_partition_count); } } else if (_enable_insert_rebalance) { - _apply_insert_rebalance(ops, insert_hashes, block->bytes()); + RETURN_IF_ERROR(_apply_insert_rebalance(ops, insert_hashes, block->bytes())); } } @@ -276,14 +266,14 @@ Status MergePartitioner::clone(RuntimeState* state, std::unique_ptr& ops, - std::vector& insert_hashes, - size_t block_bytes) const { +Status MergePartitioner::_apply_insert_rebalance(const std::vector& ops, + std::vector& insert_hashes, + size_t block_bytes) const { if (!_enable_insert_rebalance || _insert_writer_assigner == nullptr) { - return; + return Status::OK(); } if (insert_hashes.empty() || _insert_partition_count == 0) { - return; + return Status::OK(); } std::vector mask(ops.size(), 0); for (size_t i = 0; i < ops.size(); ++i) { @@ -291,7 +281,8 @@ void MergePartitioner::_apply_insert_rebalance(const std::vector& ops, mask[i] = 1; } } - _insert_writer_assigner->assign(insert_hashes, &mask, ops.size(), block_bytes, insert_hashes); + return _insert_writer_assigner->assign(insert_hashes, &mask, ops.size(), block_bytes, + insert_hashes); } void MergePartitioner::_init_insert_scaling(RuntimeState* state) { @@ -324,10 +315,10 @@ void MergePartitioner::_init_insert_scaling(RuntimeState* state) { } int task_num = state == nullptr ? 0 : state->task_num(); - int64_t min_partition_threshold = scale_threshold_by_task( + int64_t min_partition_threshold = scale_writer_threshold_by_task( config::table_sink_partition_write_min_partition_data_processed_rebalance_threshold, task_num); - int64_t min_data_threshold = scale_threshold_by_task( + int64_t min_data_threshold = scale_writer_threshold_by_task( config::table_sink_partition_write_min_data_processed_rebalance_threshold, task_num); _insert_writer_assigner = std::make_unique( diff --git a/be/src/format/transformer/merge_partitioner.h b/be/src/format/transformer/merge_partitioner.h index 14619c8eca2f3e..3cc7420344bfcf 100644 --- a/be/src/format/transformer/merge_partitioner.h +++ b/be/src/format/transformer/merge_partitioner.h @@ -22,7 +22,7 @@ #include #include "exec/partitioner/partitioner.h" -#include "format/transformer/writer_assigner.h" +#include "exec/partitioner/writer_assigner.h" namespace doris { @@ -40,8 +40,8 @@ class MergePartitioner final : public PartitionerBase { Status clone(RuntimeState* state, std::unique_ptr& partitioner) override; private: - void _apply_insert_rebalance(const std::vector& ops, - std::vector& insert_hashes, size_t block_bytes) const; + Status _apply_insert_rebalance(const std::vector& ops, + std::vector& insert_hashes, size_t block_bytes) const; void _init_insert_scaling(RuntimeState* state); uint32_t _next_rr_channel() const; Status _clone_expr_ctxs(RuntimeState* state, const VExprContextSPtrs& src, diff --git a/be/src/format/transformer/writer_assigner.h b/be/src/format/transformer/writer_assigner.h deleted file mode 100644 index 4c22862178b8a2..00000000000000 --- a/be/src/format/transformer/writer_assigner.h +++ /dev/null @@ -1,125 +0,0 @@ -// 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. - -#pragma once - -#include -#include -#include -#include - -#include "exec/connector/skewed_partition_rebalancer.h" - -namespace doris { - -class WriterAssigner { -public: - virtual ~WriterAssigner() = default; - - virtual void assign(const std::vector& partition_ids, - const std::vector* mask, size_t rows, size_t block_bytes, - std::vector& writer_ids) = 0; -}; - -class IdentityWriterAssigner final : public WriterAssigner { -public: - void assign(const std::vector& partition_ids, const std::vector* mask, - size_t rows, size_t /*block_bytes*/, std::vector& writer_ids) override { - if (rows == 0) { - return; - } - if (writer_ids.size() != rows && &writer_ids != &partition_ids) { - writer_ids.resize(rows); - } - if (mask == nullptr) { - for (size_t i = 0; i < rows; ++i) { - writer_ids[i] = partition_ids[i]; - } - return; - } - for (size_t i = 0; i < rows; ++i) { - if ((*mask)[i] == 0) { - continue; - } - writer_ids[i] = partition_ids[i]; - } - } -}; - -class SkewedWriterAssigner final : public WriterAssigner { -public: - SkewedWriterAssigner(int partition_count, int task_count, int task_bucket_count, - long min_partition_data_processed_rebalance_threshold, - long min_data_processed_rebalance_threshold) - : _rebalancer(partition_count, task_count, task_bucket_count, - min_partition_data_processed_rebalance_threshold, - min_data_processed_rebalance_threshold), - _partition_row_counts(partition_count, 0), - _partition_writer_ids(partition_count, -1), - _partition_writer_indexes(partition_count, 0) {} - - void assign(const std::vector& partition_ids, const std::vector* mask, - size_t rows, size_t block_bytes, std::vector& writer_ids) override { - if (rows == 0 || _partition_row_counts.empty()) { - return; - } - if (writer_ids.size() != rows && &writer_ids != &partition_ids) { - writer_ids.resize(rows); - } - - std::fill(_partition_row_counts.begin(), _partition_row_counts.end(), 0); - std::fill(_partition_writer_ids.begin(), _partition_writer_ids.end(), -1); - _rebalancer.rebalance(); - - const size_t partition_count = _partition_row_counts.size(); - for (size_t i = 0; i < rows; ++i) { - if (mask != nullptr && (*mask)[i] == 0) { - continue; - } - const uint32_t partition_id = partition_ids[i]; - if (partition_id >= partition_count) { - continue; - } - _partition_row_counts[partition_id] += 1; - int writer_id = _partition_writer_ids[partition_id]; - if (writer_id == -1) { - writer_id = _get_next_writer_id(partition_id); - _partition_writer_ids[partition_id] = writer_id; - } - writer_ids[i] = static_cast(writer_id); - } - - for (size_t i = 0; i < partition_count; ++i) { - if (_partition_row_counts[i] > 0) { - _rebalancer.add_partition_row_count(static_cast(i), _partition_row_counts[i]); - } - } - _rebalancer.add_data_processed(static_cast(block_bytes)); - } - -private: - int _get_next_writer_id(uint32_t partition_id) { - return _rebalancer.get_task_id(partition_id, _partition_writer_indexes[partition_id]++); - } - - SkewedPartitionRebalancer _rebalancer; - std::vector _partition_row_counts; - std::vector _partition_writer_ids; - std::vector _partition_writer_indexes; -}; - -} // namespace doris diff --git a/be/src/runtime/runtime_state.cpp b/be/src/runtime/runtime_state.cpp index 50740802417aeb..049d6b165c9012 100644 --- a/be/src/runtime/runtime_state.cpp +++ b/be/src/runtime/runtime_state.cpp @@ -72,18 +72,37 @@ Status RuntimeState::add_iceberg_commit_datas(TIcebergCommitData iceberg_commit_ thrift_limit > report_envelope_headroom ? thrift_limit - report_envelope_headroom : 0; std::lock_guard budget_lock(_external_file_report_state->mutex); // Parallel task states share this budget because FE receives their vectors in one fragment report. - if (_external_file_report_state->iceberg_serialized_bytes + serialized_size + sizeof(uint32_t) > + if (_external_file_report_state->serialized_commit_bytes + serialized_size + sizeof(uint32_t) > commit_data_limit) { return Status::InternalError( "Iceberg commit metadata exceeds the Thrift report limit; reduce output file " "count"); } std::lock_guard data_lock(_iceberg_commit_datas_mutex); - _external_file_report_state->iceberg_serialized_bytes += serialized_size + sizeof(uint32_t); + _external_file_report_state->serialized_commit_bytes += serialized_size + sizeof(uint32_t); _iceberg_commit_datas.emplace_back(std::move(iceberg_commit_data)); return Status::OK(); } +Status RuntimeState::add_connector_commit_data(std::string commit_data) { + constexpr size_t report_envelope_headroom = 1024 * 1024; + const size_t thrift_limit = coordinator_thrift_message_limit(); + const size_t commit_data_limit = + thrift_limit > report_envelope_headroom ? thrift_limit - report_envelope_headroom : 0; + std::lock_guard budget_lock(_external_file_report_state->mutex); + if (_external_file_report_state->serialized_commit_bytes + commit_data.size() + + sizeof(uint32_t) > + commit_data_limit) { + return Status::InternalError( + "Connector commit metadata exceeds the Thrift report limit; reduce commit " + "metadata size"); + } + std::lock_guard data_lock(_connector_commit_data_mutex); + _external_file_report_state->serialized_commit_bytes += commit_data.size() + sizeof(uint32_t); + _connector_commit_data.emplace_back(std::move(commit_data)); + return Status::OK(); +} + size_t RuntimeState::coordinator_thrift_message_limit() const { int32_t effective_thrift_limit = std::max(config::thrift_max_message_size, 0); if (_query_options.__isset.coordinator_thrift_max_message_size && @@ -115,6 +134,10 @@ void RuntimeState::append_external_file_commit_data(TReportExecStatusParams* par params->mc_commit_datas.insert(params->mc_commit_datas.end(), commit_datas.begin(), commit_datas.end()); } + append_connector_commit_data(¶ms->connector_commit_data); + if (!params->connector_commit_data.empty()) { + params->__isset.connector_commit_data = true; + } } void RuntimeState::add_rejected_external_file_report_cleanup(std::function cleanup) { diff --git a/be/src/runtime/runtime_state.h b/be/src/runtime/runtime_state.h index bd9b849ba7a4de..def24619e98973 100644 --- a/be/src/runtime/runtime_state.h +++ b/be/src/runtime/runtime_state.h @@ -84,7 +84,7 @@ class ExternalFileReportState { private: std::mutex mutex; - size_t iceberg_serialized_bytes = 0; + size_t serialized_commit_bytes = 0; bool ownership_may_have_transferred = false; std::vector> rejected_report_cleanups; }; @@ -547,6 +547,13 @@ class RuntimeState { Status add_iceberg_commit_datas(TIcebergCommitData iceberg_commit_data); + Status add_connector_commit_data(std::string commit_data); + + void append_connector_commit_data(std::vector* output) const { + std::lock_guard lock(_connector_commit_data_mutex); + output->insert(output->end(), _connector_commit_data.begin(), _connector_commit_data.end()); + } + size_t coordinator_thrift_message_limit() const; void append_external_file_commit_data(TReportExecStatusParams* params, bool final_report) const; @@ -1012,6 +1019,9 @@ class RuntimeState { mutable std::mutex _mc_commit_datas_mutex; std::vector _mc_commit_datas; + mutable std::mutex _connector_commit_data_mutex; + std::vector _connector_commit_data; + std::vector> _op_id_to_local_state; std::unique_ptr _sink_local_state; diff --git a/be/test/core/data_type_serde/data_type_serde_arrow_test.cpp b/be/test/core/data_type_serde/data_type_serde_arrow_test.cpp index 6f87a555ed7d73..aaafe9c30c1c8d 100644 --- a/be/test/core/data_type_serde/data_type_serde_arrow_test.cpp +++ b/be/test/core/data_type_serde/data_type_serde_arrow_test.cpp @@ -80,6 +80,7 @@ #include "core/data_type/data_type_timestamp_ns.h" #include "core/data_type/data_type_timestamptz.h" #include "core/data_type/data_type_varbinary.h" +#include "core/data_type/data_type_variant_v2.h" #include "core/data_type/define_primitive_type.h" #include "core/field.h" #include "core/types.h" diff --git a/be/test/core/data_type_serde/data_type_variant_v2_serde_output_test.cpp b/be/test/core/data_type_serde/data_type_variant_v2_serde_output_test.cpp index 9f5e6054a1fba7..5aa88ab50c6500 100644 --- a/be/test/core/data_type_serde/data_type_variant_v2_serde_output_test.cpp +++ b/be/test/core/data_type_serde/data_type_variant_v2_serde_output_test.cpp @@ -15,7 +15,9 @@ // specific language governing permissions and limitations // under the License. +#include #include +#include #include #include @@ -198,6 +200,50 @@ std::vector> orc_values(const DataTypeVariantV2SerDe& return result; } +std::shared_ptr binary_variant_arrow_type() { + return arrow::struct_({arrow::field("value", arrow::binary(), false), + arrow::field("metadata", arrow::binary(), false)}); +} + +std::unique_ptr binary_variant_arrow_builder() { + return std::make_unique( + binary_variant_arrow_type(), arrow::default_memory_pool(), + std::vector> { + std::make_shared(arrow::default_memory_pool()), + std::make_shared(arrow::default_memory_pool())}); +} + +void expect_binary_variant_bytes(const DataTypeVariantV2SerDe& serde, const IColumn& column, + const ColumnVariantV2& encoded, + const NullMap* null_map = nullptr) { + auto builder = binary_variant_arrow_builder(); + const Status status = serde.write_column_to_arrow(column, null_map, builder.get(), 0, + column.size(), cctz::utc_time_zone()); + ASSERT_TRUE(status.ok()) << status; + + std::shared_ptr output; + ASSERT_TRUE(builder->Finish(&output).ok()); + const auto& array = assert_cast(*output); + const auto& values = assert_cast(*array.field(0)); + const auto& metadata = assert_cast(*array.field(1)); + ASSERT_EQ(array.length(), static_cast(column.size())); + const auto view = encoded.read_view(); + for (size_t row = 0; row < column.size(); ++row) { + const bool expected_null = null_map != nullptr && (*null_map)[row] != 0; + EXPECT_EQ(array.IsNull(row), expected_null); + if (expected_null) { + continue; + } + const VariantRef expected = view.value_at(row); + const auto actual_value = values.GetView(row); + const auto actual_metadata = metadata.GetView(row); + EXPECT_EQ(std::string_view(actual_value.data(), actual_value.size()), + std::string_view(expected.value.data, expected.value.size)); + EXPECT_EQ(std::string_view(actual_metadata.data(), actual_metadata.size()), + std::string_view(expected.metadata.data, expected.metadata.size)); + } +} + // NOLINTNEXTLINE(readability-function-cognitive-complexity) -- GTest macros inflate the matrix. void expect_text_surfaces(const DataTypeVariantV2SerDe& serde, const IColumn& encoded, const ColumnVariantV2& typed, @@ -397,4 +443,35 @@ TEST(DataTypeVariantV2SerdeOutputTest, ConstNullableAndOuterMasksPreserveBoundar EXPECT_TRUE(invalid_dates->is_typed()); } +TEST(DataTypeVariantV2SerdeOutputTest, BinaryStructPreservesEncodedAndTypedBytesAndOuterNulls) { + DataTypeVariantV2SerDe serde; + auto documents = encoded_json({R"({"a":[1,null,"x"]})", R"({"hidden":true})", "null"}); + NullMap mask {0, 1, 0}; + expect_binary_variant_bytes(serde, *documents, *documents, &mask); + + auto typed = typed_strings( + {std::string_view("plain"), std::nullopt, std::string_view(R"({"text":"value"})")}); + ColumnPtr encoded = encoded_copy(*typed); + expect_binary_variant_bytes(serde, *typed, assert_cast(*encoded)); + EXPECT_TRUE(typed->is_typed()); +} + +TEST(DataTypeVariantV2SerdeOutputTest, BinaryStructRejectsUnsupportedPrimitive) { + DataTypeVariantV2SerDe serde; + VariantBatchBuilder builder(VariantBatchBuilder::ReserveHint {.rows = 1}); + auto row = builder.begin_row(); + row.add_time_ntz_micros(1'500'000); + row.finish(); + auto encoded = ColumnVariantV2::create(); + encoded->insert_encoded_batch(builder.finish_batch()); + auto arrow_builder = binary_variant_arrow_builder(); + const Status status = serde.write_column_to_arrow(*encoded, nullptr, arrow_builder.get(), 0, + encoded->size(), cctz::utc_time_zone()); + EXPECT_EQ(status.code(), ErrorCode::NOT_IMPLEMENTED_ERROR); + EXPECT_NE(status.to_string().find( + "Binary Variant V2 Arrow encoding does not support primitive id 17"), + std::string::npos); + EXPECT_EQ(arrow_builder->length(), 0); +} + } // namespace doris diff --git a/be/test/exec/operator/exchange_sink_operator_test.cpp b/be/test/exec/operator/exchange_sink_operator_test.cpp index 2a65907c58ff23..77ec95fcbc4de2 100644 --- a/be/test/exec/operator/exchange_sink_operator_test.cpp +++ b/be/test/exec/operator/exchange_sink_operator_test.cpp @@ -20,6 +20,7 @@ #include #include +#include #include #include "core/block/block.h" @@ -152,6 +153,29 @@ TEST(ExchangeSinkOperatorTest, test_all_remote) { {.is_local = false, .fragment_instance_id = create_TUniqueId(1, 5)}}); } +TEST(ExchangeSinkOperatorTest, shared_writer_scaling_state_is_synchronized) { + auto [op, ctx, mock_channel] = create_exchange_sink( + {{.is_local = true, .fragment_instance_id = create_TUniqueId(1, 1)}}); + constexpr size_t thread_count = 8; + constexpr size_t updates_per_thread = 100; + std::vector threads; + threads.reserve(thread_count); + for (size_t i = 0; i < thread_count; ++i) { + threads.emplace_back([&] { + for (size_t update = 0; update < updates_per_thread; ++update) { + op->update_writer_scaling_for_test(1, 1); + } + }); + } + for (auto& thread : threads) { + thread.join(); + } + + auto [data_processed, writer_count] = op->writer_scaling_state_for_test(); + EXPECT_EQ(data_processed, thread_count * updates_per_thread); + EXPECT_EQ(writer_count, 1); +} + TEST(ExchangeSinkOperatorTest, test_some_api) { auto [op, ctx, mock_channel] = create_exchange_sink( {{.is_local = true, .fragment_instance_id = create_TUniqueId(1, 1)}, diff --git a/be/test/exec/partitioner/external_partition_function_factory_test.cpp b/be/test/exec/partitioner/external_partition_function_factory_test.cpp new file mode 100644 index 00000000000000..f7ec064bb9c80c --- /dev/null +++ b/be/test/exec/partitioner/external_partition_function_factory_test.cpp @@ -0,0 +1,65 @@ +// 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/partitioner/external/external_partition_function_factory.h" + +#include + +#include +#include +#include + +namespace doris { + +TEST(ExternalPartitionFunctionFactoryTest, CreatesGenericDirectHash) { + TExternalTableSinkHashPartitionInfo info; + info.__set_partition_function("direct_hash"); + info.__set_writer_assignment(TExternalTableSinkWriterAssignment::IDENTITY); + std::unique_ptr function; + + ASSERT_TRUE(create_external_partition_function(info, 3, ShuffleHashMethod::CRC32, {}, &function) + .ok()); + ASSERT_NE(function, nullptr); + EXPECT_EQ(function->partition_count(), 3); +} + +TEST(ExternalPartitionFunctionFactoryTest, RejectsUnknownFunction) { + TExternalTableSinkHashPartitionInfo info; + info.__set_partition_function("connector_owned_function"); + info.__set_writer_assignment(TExternalTableSinkWriterAssignment::IDENTITY); + std::unique_ptr function; + + Status status = + create_external_partition_function(info, 3, ShuffleHashMethod::CRC32, {}, &function); + EXPECT_FALSE(status.ok()); + EXPECT_NE(status.to_string().find("Unsupported external sink partition function"), + std::string::npos); +} + +TEST(ExternalPartitionFunctionFactoryTest, DirectHashRejectsOpaqueOptions) { + TExternalTableSinkHashPartitionInfo info; + info.__set_partition_function("direct_hash"); + info.__set_partition_function_options({{"unexpected", "value"}}); + info.__set_writer_assignment(TExternalTableSinkWriterAssignment::IDENTITY); + std::unique_ptr function; + + Status status = + create_external_partition_function(info, 3, ShuffleHashMethod::CRC32, {}, &function); + EXPECT_FALSE(status.ok()); + EXPECT_NE(status.to_string().find("does not accept options"), std::string::npos); +} + +} // namespace doris diff --git a/be/test/exec/partitioner/writer_assigner_test.cpp b/be/test/exec/partitioner/writer_assigner_test.cpp new file mode 100644 index 00000000000000..04883d2ed1543f --- /dev/null +++ b/be/test/exec/partitioner/writer_assigner_test.cpp @@ -0,0 +1,55 @@ +// 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/partitioner/writer_assigner.h" + +#include + +#include +#include + +namespace doris { + +TEST(WriterAssignerTest, IdentityPreservesLogicalPartition) { + IdentityWriterAssigner assigner(3); + std::vector partition_ids {2, 0, 1, 2}; + std::vector writer_ids; + + ASSERT_TRUE(assigner.assign(partition_ids, nullptr, partition_ids.size(), 64, writer_ids).ok()); + EXPECT_EQ(partition_ids, writer_ids); +} + +TEST(WriterAssignerTest, IdentityRejectsInvalidLogicalPartition) { + IdentityWriterAssigner assigner(2); + std::vector partition_ids {0, 2}; + std::vector writer_ids; + + Status status = assigner.assign(partition_ids, nullptr, partition_ids.size(), 64, writer_ids); + EXPECT_FALSE(status.ok()); + EXPECT_NE(status.to_string().find("exceeds writer count"), std::string::npos); +} + +TEST(WriterAssignerTest, SkewedRejectsInvalidLogicalPartition) { + SkewedWriterAssigner assigner(4, 2, 1, 1, 1); + std::vector partition_ids {0, 4}; + std::vector writer_ids; + + Status status = assigner.assign(partition_ids, nullptr, partition_ids.size(), 64, writer_ids); + EXPECT_FALSE(status.ok()); + EXPECT_NE(status.to_string().find("exceeds partition count"), std::string::npos); +} + +} // namespace doris diff --git a/be/test/exprs/function/cast/cast_variant_v2_from_test.cpp b/be/test/exprs/function/cast/cast_variant_v2_from_test.cpp index 6604d59f976440..d00d3587c6856a 100644 --- a/be/test/exprs/function/cast/cast_variant_v2_from_test.cpp +++ b/be/test/exprs/function/cast/cast_variant_v2_from_test.cpp @@ -32,6 +32,7 @@ #include "core/data_type/data_type_decimal.h" #include "core/data_type/data_type_ipv6.h" #include "core/data_type/data_type_jsonb.h" +#include "core/data_type/data_type_nothing.h" #include "core/data_type/data_type_nullable.h" #include "core/data_type/data_type_number.h" #include "core/data_type/data_type_string.h" @@ -780,6 +781,29 @@ TEST(CastVariantV2FromTest, NestedArrayRoundTripPreservesNullAndEmptyArray) { EXPECT_EQ(assert_cast(values.get_nested_column()).get_data()[0], 1); } +TEST(CastVariantV2FromTest, NullOnlyArrayEncodesNonEmptyElements) { + auto array_type = std::make_shared(std::make_shared()); + MutableColumnPtr source = array_type->create_column(); + Array values {Field::create_field(Null()), Field::create_field(Null())}; + source->insert(Field::create_field(std::move(values))); + + auto variant_type = std::make_shared(); + Block block {{source->get_ptr(), array_type, "source"}, + {variant_type->create_column(), variant_type, "result"}}; + RuntimeState state; + auto context = FunctionContext::create_context(&state, {}, {}); + Status status = + create_cast_to_variant_v2_wrapper(array_type)(context.get(), block, {0}, 1, 1, nullptr); + ASSERT_TRUE(status.ok()) << status; + + VariantRef encoded = + assert_cast(*block.get_by_position(1).column).get_value_ref(0); + ASSERT_EQ(encoded.basic_type(), VariantBasicType::ARRAY); + ASSERT_EQ(encoded.num_elements(), 2); + EXPECT_TRUE(encoded.array_at(0).is_null()); + EXPECT_TRUE(encoded.array_at(1).is_null()); +} + TEST(CastVariantV2FromTest, DecimalScale38CastsAndScale39IsRejectedAtEncodingBoundary) { VariantBatchBuilder builder(VariantBatchBuilder::ReserveHint {.rows = 1}); auto row = builder.begin_row(); diff --git a/be/test/runtime/runtime_state_block_budget_test.cpp b/be/test/runtime/runtime_state_block_budget_test.cpp index 5a384378ec382b..9e7dca7fdbe29c 100644 --- a/be/test/runtime/runtime_state_block_budget_test.cpp +++ b/be/test/runtime/runtime_state_block_budget_test.cpp @@ -60,6 +60,28 @@ TEST(RuntimeStateIcebergCommitDataTest, SharesTheReportBudgetAcrossParallelTasks EXPECT_FALSE(second_status.ok()); } +TEST(RuntimeStateIcebergCommitDataTest, SharesTheReportBudgetWithOpaqueConnectorData) { + RuntimeState iceberg_state; + RuntimeState connector_state; + auto budget = std::make_shared(); + iceberg_state.set_external_file_report_state(budget); + connector_state.set_external_file_report_state(budget); + const int32_t saved_limit = config::thrift_max_message_size; + config::thrift_max_message_size = 1024 * 1024 + 512; + TIcebergCommitData iceberg_data; + iceberg_data.__set_file_path(std::string(300, 'x')); + + Status first_status = iceberg_state.add_iceberg_commit_datas(iceberg_data); + Status second_status = connector_state.add_connector_commit_data(std::string(300, 'y')); + + config::thrift_max_message_size = saved_limit; + EXPECT_TRUE(first_status.ok()) << first_status; + EXPECT_FALSE(second_status.ok()); + std::vector collected; + connector_state.append_connector_commit_data(&collected); + EXPECT_TRUE(collected.empty()); +} + TEST(RuntimeStateIcebergCommitDataTest, UsesTheSmallerCoordinatorThriftLimit) { RuntimeState state; const int32_t saved_limit = config::thrift_max_message_size; @@ -91,6 +113,7 @@ TEST(RuntimeStateIcebergCommitDataTest, PeriodicReportOmitsExternalCommitData) { ASSERT_TRUE(state.add_iceberg_commit_datas(iceberg_data).ok()); TMCCommitData mc_data; state.add_mc_commit_datas(mc_data); + ASSERT_TRUE(state.add_connector_commit_data("opaque-fragment").ok()); TReportExecStatusParams periodic_params; state.append_external_file_commit_data(&periodic_params, false); @@ -98,12 +121,16 @@ TEST(RuntimeStateIcebergCommitDataTest, PeriodicReportOmitsExternalCommitData) { EXPECT_FALSE(periodic_params.__isset.hive_partition_updates); EXPECT_FALSE(periodic_params.__isset.iceberg_commit_datas); EXPECT_FALSE(periodic_params.__isset.mc_commit_datas); + EXPECT_FALSE(periodic_params.__isset.connector_commit_data); TReportExecStatusParams final_params; state.append_external_file_commit_data(&final_params, true); EXPECT_TRUE(final_params.__isset.hive_partition_updates); EXPECT_TRUE(final_params.__isset.iceberg_commit_datas); EXPECT_TRUE(final_params.__isset.mc_commit_datas); + ASSERT_TRUE(final_params.__isset.connector_commit_data); + ASSERT_EQ(1, final_params.connector_commit_data.size()); + EXPECT_EQ("opaque-fragment", final_params.connector_commit_data[0]); } TEST(RuntimeStateIcebergCommitDataTest, RetainsFileCleanupUntilReportAcknowledgement) { diff --git a/fe/fe-catalog/src/main/java/org/apache/doris/analysis/ShortCircuitFunctionCallExpr.java b/fe/fe-catalog/src/main/java/org/apache/doris/analysis/ShortCircuitFunctionCallExpr.java new file mode 100644 index 00000000000000..ef254ab518d32e --- /dev/null +++ b/fe/fe-catalog/src/main/java/org/apache/doris/analysis/ShortCircuitFunctionCallExpr.java @@ -0,0 +1,39 @@ +// 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. + +package org.apache.doris.analysis; + +import org.apache.doris.catalog.Function; + +/** Function call whose unselected arguments must not be evaluated. */ +public final class ShortCircuitFunctionCallExpr extends FunctionCallExpr { + + /** Create a function call with mandatory lazy branch evaluation. */ + public ShortCircuitFunctionCallExpr( + Function function, FunctionParams functionParams, boolean nullable) { + super(function, functionParams, nullable); + } + + private ShortCircuitFunctionCallExpr(ShortCircuitFunctionCallExpr other) { + super(other); + } + + @Override + public Expr clone() { + return new ShortCircuitFunctionCallExpr(this); + } +} diff --git a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergWritePlanProvider.java b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergWritePlanProvider.java index 80597f6610727c..83483dc83b8478 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergWritePlanProvider.java +++ b/fe/fe-connector/fe-connector-iceberg/src/main/java/org/apache/doris/connector/iceberg/IcebergWritePlanProvider.java @@ -29,6 +29,7 @@ import org.apache.doris.connector.spi.handle.ConnectorTransaction; import org.apache.doris.connector.spi.handle.ConnectorWriteHandle; import org.apache.doris.connector.spi.handle.WriteOperation; +import org.apache.doris.connector.spi.write.ConnectorRowChangeStyle; import org.apache.doris.connector.spi.write.ConnectorSinkPlan; import org.apache.doris.connector.spi.write.ConnectorWritePartitionField; import org.apache.doris.connector.spi.write.ConnectorWritePartitionSpec; @@ -105,6 +106,11 @@ */ public class IcebergWritePlanProvider implements ConnectorWritePlanProvider { + @Override + public ConnectorRowChangeStyle getRowChangeStyle() { + return ConnectorRowChangeStyle.POSITION_DELETE; + } + private static final int SUPPORT_NESTED_PARTITION_WRITE_EXEC_VERSION = 12; // Legacy IcebergUtils compression-codec property keys (connector-local copies; iceberg SDK has no @@ -119,6 +125,11 @@ public class IcebergWritePlanProvider implements ConnectorWritePlanProvider { // drift on either side turns one of the two tests red. private static final String DORIS_ICEBERG_ROWID_COL = "__DORIS_ICEBERG_ROWID_COL__"; + private static final Set ROW_LEVEL_WRITE_CONSTRAINT_EXCLUDED_COLUMNS = + Collections.unmodifiableSet(new HashSet<>(Arrays.asList( + DORIS_ICEBERG_ROWID_COL, "$file_path", "$row_position", + "$partition_spec_id", "$partition_data"))); + // The single request-scoped synthetic write column iceberg declares: the row-id STRUCT carrying the // per-row write metadata (file_path / row_position / partition_spec_id / partition_data). Same for // every iceberg table regardless of format/partitioning, so it is a shared immutable instance. @@ -643,10 +654,29 @@ public List getSyntheticWriteColumns(ConnectorSession session, return SYNTHETIC_WRITE_COLUMNS; } + @Override + public Set getRowLevelWriteConstraintExcludedColumns() { + return ROW_LEVEL_WRITE_CONSTRAINT_EXCLUDED_COLUMNS; + } + + @Override + public String getRowLevelDmlLabelPrefix(WriteOperation operation) { + switch (operation) { + case DELETE: + return "iceberg_delete"; + case UPDATE: + return "iceberg_update_merge"; + case MERGE: + return "iceberg_merge_into"; + default: + throw new DorisConnectorException("Unsupported Iceberg row-level operation: " + operation); + } + } + @Override public Set supportedOperations() { return EnumSet.of(WriteOperation.INSERT, WriteOperation.OVERWRITE, - WriteOperation.DELETE, WriteOperation.MERGE, WriteOperation.REWRITE); + WriteOperation.DELETE, WriteOperation.UPDATE, WriteOperation.MERGE, WriteOperation.REWRITE); } @Override diff --git a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergConnectorTest.java b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergConnectorTest.java index 78b5ec6e3af167..ee671040606e22 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergConnectorTest.java +++ b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergConnectorTest.java @@ -365,7 +365,7 @@ public void declaredWriteCapabilitiesMatchAndPassContractValidator() throws Exce Assertions.assertNotNull(writeProvider, "iceberg connector must expose a write plan provider"); Assertions.assertEquals( EnumSet.of(WriteOperation.INSERT, WriteOperation.OVERWRITE, WriteOperation.DELETE, - WriteOperation.MERGE, WriteOperation.REWRITE), + WriteOperation.UPDATE, WriteOperation.MERGE, WriteOperation.REWRITE), writeProvider.supportedOperations()); Assertions.assertTrue(writeProvider.supportsWriteBranch()); Assertions.assertTrue(writeProvider.requiresParallelWrite()); diff --git a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergWritePlanProviderTest.java b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergWritePlanProviderTest.java index ad03f8a0a2039b..4d5040992e4477 100644 --- a/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergWritePlanProviderTest.java +++ b/fe/fe-connector/fe-connector-iceberg/src/test/java/org/apache/doris/connector/iceberg/IcebergWritePlanProviderTest.java @@ -28,6 +28,7 @@ import org.apache.doris.connector.spi.handle.ConnectorWriteHandle; import org.apache.doris.connector.spi.handle.WriteOperation; import org.apache.doris.connector.spi.mvcc.ConnectorMvccSnapshot; +import org.apache.doris.connector.spi.write.ConnectorRowChangeStyle; import org.apache.doris.connector.spi.write.ConnectorSinkPlan; import org.apache.doris.connector.spi.write.ConnectorWritePartitionField; import org.apache.doris.connector.spi.write.ConnectorWritePartitionSpec; @@ -1867,7 +1868,9 @@ public void declaresFullWriteOperationSet() { IcebergWritePlanProvider provider = providerFor(unpartitionedUnsortedTable(freshCatalog()), contextWithStorage()); Assertions.assertEquals(EnumSet.of(WriteOperation.INSERT, WriteOperation.OVERWRITE, - WriteOperation.DELETE, WriteOperation.MERGE, WriteOperation.REWRITE), provider.supportedOperations()); + WriteOperation.DELETE, WriteOperation.UPDATE, WriteOperation.MERGE, WriteOperation.REWRITE), + provider.supportedOperations()); + Assertions.assertEquals(ConnectorRowChangeStyle.POSITION_DELETE, provider.getRowChangeStyle()); Assertions.assertTrue(provider.supportsWriteBranch()); Assertions.assertTrue(provider.requiresParallelWrite()); Assertions.assertTrue(provider.requiresFullSchemaWriteOrder()); diff --git a/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/write/ConnectorChangelogMode.java b/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/write/ConnectorChangelogMode.java new file mode 100644 index 00000000000000..9652f4b7724a30 --- /dev/null +++ b/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/write/ConnectorChangelogMode.java @@ -0,0 +1,58 @@ +// 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. + +package org.apache.doris.connector.spi.write; + +import java.util.Objects; + +/** Connector-owned encoding of row operations in a changelog write. */ +public final class ConnectorChangelogMode { + private final String operationColumnName; + private final byte insertValue; + private final byte updateValue; + private final byte deleteValue; + + public ConnectorChangelogMode(String operationColumnName, + byte insertValue, byte updateValue, byte deleteValue) { + this.operationColumnName = Objects.requireNonNull(operationColumnName, "operationColumnName"); + if (operationColumnName.isEmpty()) { + throw new IllegalArgumentException("Changelog operation column name must not be empty"); + } + if (insertValue == updateValue || insertValue == deleteValue || updateValue == deleteValue) { + throw new IllegalArgumentException("Changelog operation values must be distinct"); + } + this.insertValue = insertValue; + this.updateValue = updateValue; + this.deleteValue = deleteValue; + } + + public String getOperationColumnName() { + return operationColumnName; + } + + public byte getInsertValue() { + return insertValue; + } + + public byte getUpdateValue() { + return updateValue; + } + + public byte getDeleteValue() { + return deleteValue; + } +} diff --git a/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/write/ConnectorRowChangeStyle.java b/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/write/ConnectorRowChangeStyle.java new file mode 100644 index 00000000000000..4ad896557c1ece --- /dev/null +++ b/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/write/ConnectorRowChangeStyle.java @@ -0,0 +1,28 @@ +// 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. + +package org.apache.doris.connector.spi.write; + +/** Physical representation used by a connector for row-level changes. */ +public enum ConnectorRowChangeStyle { + /** No row-level write plan is available. */ + NONE, + /** Deletes identify positions in existing data files. */ + POSITION_DELETE, + /** Writes encode inserts, updates, and deletes as tagged rows. */ + CHANGELOG +} diff --git a/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/write/ConnectorRowLevelDmlRequest.java b/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/write/ConnectorRowLevelDmlRequest.java new file mode 100644 index 00000000000000..12e67becee9ae7 --- /dev/null +++ b/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/write/ConnectorRowLevelDmlRequest.java @@ -0,0 +1,58 @@ +// 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. + +package org.apache.doris.connector.spi.write; + +import org.apache.doris.connector.spi.handle.WriteOperation; + +import java.util.Collections; +import java.util.Set; +import java.util.TreeSet; + +/** Connector-neutral facts needed to validate a row-level DML statement. */ +public final class ConnectorRowLevelDmlRequest { + private final WriteOperation operation; + private final Set updatedColumns; + private final boolean containsUpdate; + private final boolean containsDelete; + + public ConnectorRowLevelDmlRequest(WriteOperation operation, Set updatedColumns, + boolean containsUpdate, boolean containsDelete) { + this.operation = operation; + TreeSet columns = new TreeSet<>(String.CASE_INSENSITIVE_ORDER); + columns.addAll(updatedColumns); + this.updatedColumns = Collections.unmodifiableSet(columns); + this.containsUpdate = containsUpdate; + this.containsDelete = containsDelete; + } + + public WriteOperation getOperation() { + return operation; + } + + public Set getUpdatedColumns() { + return updatedColumns; + } + + public boolean containsUpdate() { + return containsUpdate; + } + + public boolean containsDelete() { + return containsDelete; + } +} diff --git a/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/write/ConnectorWriteDistribution.java b/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/write/ConnectorWriteDistribution.java new file mode 100644 index 00000000000000..dac651f783e0b5 --- /dev/null +++ b/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/write/ConnectorWriteDistribution.java @@ -0,0 +1,130 @@ +// 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. + +package org.apache.doris.connector.spi.write; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; + +/** Describes how rows of one connector write must be distributed to sink writers. */ +public final class ConnectorWriteDistribution { + + /** Distribution behavior understood by the engine. */ + public enum Mode { + EXECUTION_ANY, + GATHER, + HASH, + EXTERNAL_UNPARTITIONED, + EXTERNAL_HASH + } + + /** Whether one ownership key is stable on one writer or may scale across writers. */ + public enum WriterAssignment { + IDENTITY, + SKEWED + } + + private final Mode mode; + private final List routeColumns; + private final String partitionFunction; + private final Map partitionFunctionOptions; + private final WriterAssignment writerAssignment; + + private ConnectorWriteDistribution(Mode mode, List routeColumns, + String partitionFunction, Map partitionFunctionOptions, + WriterAssignment writerAssignment) { + this.mode = Objects.requireNonNull(mode, "mode must not be null"); + this.routeColumns = immutableList(routeColumns); + this.partitionFunction = partitionFunction; + this.partitionFunctionOptions = immutableMap(partitionFunctionOptions); + this.writerAssignment = writerAssignment; + } + + /** Creates a distribution mode that carries no routing metadata. */ + public static ConnectorWriteDistribution simple(Mode mode) { + if (mode == Mode.HASH || mode == Mode.EXTERNAL_HASH) { + throw new IllegalArgumentException(mode + " requires routing metadata"); + } + return new ConnectorWriteDistribution(mode, Collections.emptyList(), null, + Collections.emptyMap(), null); + } + + /** Uses the engine's ordinary hash shuffle for the named columns. */ + public static ConnectorWriteDistribution hash(List routeColumns) { + return new ConnectorWriteDistribution(Mode.HASH, requireRouteColumns(routeColumns), null, + Collections.emptyMap(), null); + } + + /** + * Uses an external writer partition function registered in BE. FE treats {@code partitionFunction} and its + * options as opaque values and forwards them together with the resolved route expressions. + */ + public static ConnectorWriteDistribution externalHash(List routeColumns, + String partitionFunction, Map partitionFunctionOptions, + WriterAssignment writerAssignment) { + String function = Objects.requireNonNull(partitionFunction, + "partitionFunction must not be null"); + if (function.isEmpty()) { + throw new IllegalArgumentException("partitionFunction must not be empty"); + } + return new ConnectorWriteDistribution(Mode.EXTERNAL_HASH, + requireRouteColumns(routeColumns), function, partitionFunctionOptions, + Objects.requireNonNull(writerAssignment, "writerAssignment must not be null")); + } + + private static List requireRouteColumns(List columns) { + List result = immutableList(columns); + if (result.isEmpty()) { + throw new IllegalArgumentException("routeColumns must not be empty"); + } + return result; + } + + private static List immutableList(List values) { + return Collections.unmodifiableList(new ArrayList<>( + Objects.requireNonNull(values, "values must not be null"))); + } + + private static Map immutableMap(Map values) { + return Collections.unmodifiableMap(new LinkedHashMap<>( + Objects.requireNonNull(values, "values must not be null"))); + } + + public Mode getMode() { + return mode; + } + + public List getRouteColumns() { + return routeColumns; + } + + public String getPartitionFunction() { + return partitionFunction; + } + + public Map getPartitionFunctionOptions() { + return partitionFunctionOptions; + } + + public WriterAssignment getWriterAssignment() { + return writerAssignment; + } +} diff --git a/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/write/ConnectorWritePlanProvider.java b/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/write/ConnectorWritePlanProvider.java index 27a795f22d9236..0336a4f0f543cd 100644 --- a/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/write/ConnectorWritePlanProvider.java +++ b/fe/fe-connector/fe-connector-spi/src/main/java/org/apache/doris/connector/spi/write/ConnectorWritePlanProvider.java @@ -160,6 +160,16 @@ default ConnectorWritePartitionSpec getWritePartitioning(ConnectorSession sessio return null; } + /** + * Returns the table's write distribution, or {@code null} to retain the engine's existing generic + * distribution rules. FE Core treats an external hash function name and its options as opaque values and + * transports them to BE, but the named function must be registered in the BE build. + */ + default ConnectorWriteDistribution getWriteDistribution(ConnectorSession session, + ConnectorTableHandle tableHandle) { + return null; + } + /** * Declares the connector's synthetic write columns for the target — request-scoped hidden * columns the engine injects into {@code PluginDrivenExternalTable.getFullSchema()} while a write/DML @@ -191,14 +201,62 @@ default List getSyntheticWriteColumns(ConnectorSession session, * The write operations this provider can plan, in one place — the single source of truth for a * connector's write capability. Replaces the removed {@code ConnectorWriteOps} boolean methods and * the removed INSERT-support capability switch. Default: INSERT only (any write provider can at least - * append). A connector overrides this to add OVERWRITE / DELETE / MERGE / REWRITE. Connector-level - * (does not vary per table); per-table mode constraints stay in - * {@link org.apache.doris.connector.spi.ConnectorWriteOps#validateRowLevelDmlMode}. + * append). A connector overrides this to add OVERWRITE / DELETE / UPDATE / MERGE / REWRITE. The engine + * resolves the provider for the target table handle, so a heterogeneous connector may return a provider + * whose operation set varies by table. Per-table mode constraints stay in + * {@link org.apache.doris.connector.spi.ConnectorWriteOps#validateRowLevelDmlMode}. A provider + * advertising DELETE, UPDATE, or MERGE must also declare its {@link #getRowChangeStyle()}. */ default Set supportedOperations() { return EnumSet.of(WriteOperation.INSERT); } + /** + * Returns the physical representation of row-level changes planned by this provider. + * The engine selects this provider per table handle, so a heterogeneous catalog can + * use different row-level plans for different tables. Connectors that only insert + * rows retain {@link ConnectorRowChangeStyle#NONE}. + */ + default ConnectorRowChangeStyle getRowChangeStyle() { + return ConnectorRowChangeStyle.NONE; + } + + /** + * Returns the connector-owned operation-column encoding for {@link ConnectorRowChangeStyle#CHANGELOG}. + * A provider declaring another row-change style keeps the empty default. + */ + default Optional getChangelogMode() { + return Optional.empty(); + } + + /** Returns the target primary-key columns used to shape changelog DELETE and MERGE plans. */ + default List getRowLevelPrimaryKeyColumns(ConnectorSession session, + ConnectorTableHandle tableHandle) { + return Collections.emptyList(); + } + + /** Performs connector-specific validation before a row-level DML plan is synthesized. */ + default void validateRowLevelDml(ConnectorSession session, ConnectorTableHandle tableHandle, + ConnectorRowLevelDmlRequest request) { + // Default: no additional validation. + } + + /** + * Column names that must not participate in row-level optimistic-conflict predicates. + * Connectors use this for synthetic row identity and file-position metadata columns. + */ + default Set getRowLevelWriteConstraintExcludedColumns() { + return Collections.emptySet(); + } + + /** + * Stable transaction-label prefix for a row-level operation. The provider selected for + * the table supplies it so generic engine planning does not embed a connector name. + */ + default String getRowLevelDmlLabelPrefix(WriteOperation operation) { + return "connector_" + operation.name().toLowerCase(java.util.Locale.ROOT); + } + /** Whether this connector can write into a named table branch ({@code INSERT INTO t@branch(name)}). Default: no. */ default boolean supportsWriteBranch() { return false; diff --git a/fe/fe-connector/fe-connector-spi/src/test/java/org/apache/doris/connector/spi/ConnectorPluginSurfaceTest.java b/fe/fe-connector/fe-connector-spi/src/test/java/org/apache/doris/connector/spi/ConnectorPluginSurfaceTest.java index dff7baba361dee..dbcb9939d6b4d5 100644 --- a/fe/fe-connector/fe-connector-spi/src/test/java/org/apache/doris/connector/spi/ConnectorPluginSurfaceTest.java +++ b/fe/fe-connector/fe-connector-spi/src/test/java/org/apache/doris/connector/spi/ConnectorPluginSurfaceTest.java @@ -21,6 +21,8 @@ import org.apache.doris.connector.spi.handle.ConnectorWriteHandle; import org.apache.doris.connector.spi.scan.ConnectorScanPlanProvider; import org.apache.doris.connector.spi.scan.ScanNodePropertyKeys; +import org.apache.doris.connector.spi.write.ConnectorChangelogMode; +import org.apache.doris.connector.spi.write.ConnectorRowLevelDmlRequest; import org.apache.doris.connector.spi.write.ConnectorWritePlanProvider; import org.junit.jupiter.api.Assertions; @@ -85,9 +87,8 @@ public void connectorApiMajorTracksTheRecordedSurfaceChange() throws IOException Assertions.assertNotNull(in, "missing connector plugin API version resource"); version.load(in); } - // ConnectorSession's external-scan-reuse policy requires major 10: an API-9 FE does not - // provide the newly added interface method to an independently built connector plugin. - Assertions.assertEquals("10.0", version.getProperty("api.version")); + // This PR changes the connector SPI surface once, from major 10 to major 11. + Assertions.assertEquals("11.0", version.getProperty("api.version")); } /** Root entry points plus provider/handle types returned to connector plugins. */ @@ -102,6 +103,9 @@ public void connectorApiMajorTracksTheRecordedSurfaceChange() throws IOException org.apache.doris.connector.spi.mvcc.ConnectorMvccSnapshot.Builder.class, ConnectorScanPlanProvider.class, ConnectorWriteHandle.class, + ConnectorChangelogMode.class, + ConnectorRowLevelDmlRequest.class, + org.apache.doris.connector.spi.write.ConnectorWriteDistribution.class, ConnectorWritePlanProvider.class, org.apache.doris.extension.spi.Plugin.class, org.apache.doris.extension.spi.PluginFactory.class, @@ -109,7 +113,10 @@ public void connectorApiMajorTracksTheRecordedSurfaceChange() throws IOException /** Public enum constants linked directly by connector plugin bytecode. */ private static final List>> FROZEN_ENUM_TYPES = - Arrays.asList(ConnectorCapability.class); + Arrays.asList(ConnectorCapability.class, + org.apache.doris.connector.spi.write.ConnectorRowChangeStyle.class, + org.apache.doris.connector.spi.write.ConnectorWriteDistribution.Mode.class, + org.apache.doris.connector.spi.write.ConnectorWriteDistribution.WriterAssignment.class); @Test public void pluginApiSurfaceMatchesRecordedBaseline() throws IOException, IllegalAccessException { diff --git a/fe/fe-connector/fe-connector-spi/src/test/java/org/apache/doris/connector/spi/write/ConnectorRowLevelDmlContractTest.java b/fe/fe-connector/fe-connector-spi/src/test/java/org/apache/doris/connector/spi/write/ConnectorRowLevelDmlContractTest.java new file mode 100644 index 00000000000000..3c5413dd360e0e --- /dev/null +++ b/fe/fe-connector/fe-connector-spi/src/test/java/org/apache/doris/connector/spi/write/ConnectorRowLevelDmlContractTest.java @@ -0,0 +1,55 @@ +// 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. + +package org.apache.doris.connector.spi.write; + +import org.apache.doris.connector.spi.handle.WriteOperation; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.Arrays; +import java.util.HashSet; + +class ConnectorRowLevelDmlContractTest { + + @Test + void changelogModeOwnsOperationColumnAndValues() { + ConnectorChangelogMode mode = new ConnectorChangelogMode("row_operation", (byte) 3, (byte) 5, (byte) 7); + + Assertions.assertEquals("row_operation", mode.getOperationColumnName()); + Assertions.assertEquals(3, mode.getInsertValue()); + Assertions.assertEquals(5, mode.getUpdateValue()); + Assertions.assertEquals(7, mode.getDeleteValue()); + Assertions.assertThrows(IllegalArgumentException.class, + () -> new ConnectorChangelogMode("row_operation", (byte) 3, (byte) 3, (byte) 7)); + } + + @Test + void rowLevelRequestCopiesUpdatedColumnsCaseInsensitively() { + ConnectorRowLevelDmlRequest request = new ConnectorRowLevelDmlRequest( + WriteOperation.MERGE, new HashSet<>(Arrays.asList("Value", "value")), true, true); + + Assertions.assertEquals(WriteOperation.MERGE, request.getOperation()); + Assertions.assertEquals(1, request.getUpdatedColumns().size()); + Assertions.assertTrue(request.getUpdatedColumns().contains("VALUE")); + Assertions.assertTrue(request.containsUpdate()); + Assertions.assertTrue(request.containsDelete()); + Assertions.assertThrows(UnsupportedOperationException.class, + () -> request.getUpdatedColumns().add("another")); + } +} diff --git a/fe/fe-connector/fe-connector-spi/src/test/java/org/apache/doris/connector/spi/write/ConnectorWriteDistributionTest.java b/fe/fe-connector/fe-connector-spi/src/test/java/org/apache/doris/connector/spi/write/ConnectorWriteDistributionTest.java new file mode 100644 index 00000000000000..82c14bbbd269b3 --- /dev/null +++ b/fe/fe-connector/fe-connector-spi/src/test/java/org/apache/doris/connector/spi/write/ConnectorWriteDistributionTest.java @@ -0,0 +1,64 @@ +// 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. + +package org.apache.doris.connector.spi.write; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; + +class ConnectorWriteDistributionTest { + + @Test + void externalHashCopiesConnectorOwnedMetadata() { + List columns = new ArrayList<>(Arrays.asList("part", "id")); + Map options = new LinkedHashMap<>( + Collections.singletonMap("bucket-count", "8")); + ConnectorWriteDistribution distribution = ConnectorWriteDistribution.externalHash( + columns, "connector_bucket", options, + ConnectorWriteDistribution.WriterAssignment.IDENTITY); + + columns.clear(); + options.clear(); + + Assertions.assertEquals(ConnectorWriteDistribution.Mode.EXTERNAL_HASH, + distribution.getMode()); + Assertions.assertEquals(Arrays.asList("part", "id"), distribution.getRouteColumns()); + Assertions.assertEquals("connector_bucket", distribution.getPartitionFunction()); + Assertions.assertEquals(Collections.singletonMap("bucket-count", "8"), + distribution.getPartitionFunctionOptions()); + Assertions.assertEquals(ConnectorWriteDistribution.WriterAssignment.IDENTITY, + distribution.getWriterAssignment()); + } + + @Test + void hashModesRequireRoutingMetadata() { + Assertions.assertThrows(IllegalArgumentException.class, + () -> ConnectorWriteDistribution.simple(ConnectorWriteDistribution.Mode.EXTERNAL_HASH)); + Assertions.assertThrows(IllegalArgumentException.class, + () -> ConnectorWriteDistribution.hash(Collections.emptyList())); + Assertions.assertThrows(IllegalArgumentException.class, + () -> ConnectorWriteDistribution.externalHash(Collections.singletonList("id"), "", + Collections.emptyMap(), ConnectorWriteDistribution.WriterAssignment.SKEWED)); + } +} diff --git a/fe/fe-connector/fe-connector-spi/src/test/resources/connector-plugin-surface.txt b/fe/fe-connector/fe-connector-spi/src/test/resources/connector-plugin-surface.txt index 6e4eb94c2b29f1..10abdc500faffa 100644 --- a/fe/fe-connector/fe-connector-spi/src/test/resources/connector-plugin-surface.txt +++ b/fe/fe-connector/fe-connector-spi/src/test/resources/connector-plugin-surface.txt @@ -161,9 +161,41 @@ org.apache.doris.connector.spi.scan.ScanNodePropertyKeys#field:TEXT_PROPERTY_PRE org.apache.doris.connector.spi.scan.ScanNodePropertyKeys#field:TEXT_SERDE_LIB:java.lang.String=hive.text.serde_lib org.apache.doris.connector.spi.scan.ScanNodePropertyKeys#field:TEXT_SKIP_LINES:java.lang.String=hive.text.skip_lines org.apache.doris.connector.spi.scan.ScanNodePropertyKeys#field:TEXT_TRIM_DOUBLE_QUOTES:java.lang.String=hive.text.trim_double_quotes +org.apache.doris.connector.spi.write.ConnectorChangelogMode#getDeleteValue():byte +org.apache.doris.connector.spi.write.ConnectorChangelogMode#getInsertValue():byte +org.apache.doris.connector.spi.write.ConnectorChangelogMode#getOperationColumnName():java.lang.String +org.apache.doris.connector.spi.write.ConnectorChangelogMode#getUpdateValue():byte +org.apache.doris.connector.spi.write.ConnectorWriteDistribution#externalHash(java.util.List,java.lang.String,java.util.Map,org.apache.doris.connector.spi.write.ConnectorWriteDistribution$WriterAssignment):org.apache.doris.connector.spi.write.ConnectorWriteDistribution +org.apache.doris.connector.spi.write.ConnectorWriteDistribution#getMode():org.apache.doris.connector.spi.write.ConnectorWriteDistribution$Mode +org.apache.doris.connector.spi.write.ConnectorWriteDistribution#getPartitionFunction():java.lang.String +org.apache.doris.connector.spi.write.ConnectorWriteDistribution#getPartitionFunctionOptions():java.util.Map +org.apache.doris.connector.spi.write.ConnectorWriteDistribution#getRouteColumns():java.util.List +org.apache.doris.connector.spi.write.ConnectorWriteDistribution#getWriterAssignment():org.apache.doris.connector.spi.write.ConnectorWriteDistribution$WriterAssignment +org.apache.doris.connector.spi.write.ConnectorWriteDistribution#hash(java.util.List):org.apache.doris.connector.spi.write.ConnectorWriteDistribution +org.apache.doris.connector.spi.write.ConnectorWriteDistribution#simple(org.apache.doris.connector.spi.write.ConnectorWriteDistribution$Mode):org.apache.doris.connector.spi.write.ConnectorWriteDistribution +org.apache.doris.connector.spi.write.ConnectorWriteDistribution$Mode#enum:EXECUTION_ANY +org.apache.doris.connector.spi.write.ConnectorWriteDistribution$Mode#enum:EXTERNAL_HASH +org.apache.doris.connector.spi.write.ConnectorWriteDistribution$Mode#enum:EXTERNAL_UNPARTITIONED +org.apache.doris.connector.spi.write.ConnectorWriteDistribution$Mode#enum:GATHER +org.apache.doris.connector.spi.write.ConnectorWriteDistribution$Mode#enum:HASH +org.apache.doris.connector.spi.write.ConnectorWriteDistribution$WriterAssignment#enum:IDENTITY +org.apache.doris.connector.spi.write.ConnectorWriteDistribution$WriterAssignment#enum:SKEWED +org.apache.doris.connector.spi.write.ConnectorRowChangeStyle#enum:CHANGELOG +org.apache.doris.connector.spi.write.ConnectorRowChangeStyle#enum:NONE +org.apache.doris.connector.spi.write.ConnectorRowChangeStyle#enum:POSITION_DELETE +org.apache.doris.connector.spi.write.ConnectorRowLevelDmlRequest#containsDelete():boolean +org.apache.doris.connector.spi.write.ConnectorRowLevelDmlRequest#containsUpdate():boolean +org.apache.doris.connector.spi.write.ConnectorRowLevelDmlRequest#getOperation():org.apache.doris.connector.spi.handle.WriteOperation +org.apache.doris.connector.spi.write.ConnectorRowLevelDmlRequest#getUpdatedColumns():java.util.Set org.apache.doris.connector.spi.write.ConnectorWritePlanProvider#appendExplainInfo(java.lang.StringBuilder,java.lang.String,org.apache.doris.connector.spi.ConnectorSession,org.apache.doris.connector.spi.handle.ConnectorWriteHandle):void +org.apache.doris.connector.spi.write.ConnectorWritePlanProvider#getChangelogMode():java.util.Optional +org.apache.doris.connector.spi.write.ConnectorWritePlanProvider#getRowChangeStyle():org.apache.doris.connector.spi.write.ConnectorRowChangeStyle +org.apache.doris.connector.spi.write.ConnectorWritePlanProvider#getRowLevelDmlLabelPrefix(org.apache.doris.connector.spi.handle.WriteOperation):java.lang.String +org.apache.doris.connector.spi.write.ConnectorWritePlanProvider#getRowLevelPrimaryKeyColumns(org.apache.doris.connector.spi.ConnectorSession,org.apache.doris.connector.spi.handle.ConnectorTableHandle):java.util.List +org.apache.doris.connector.spi.write.ConnectorWritePlanProvider#getRowLevelWriteConstraintExcludedColumns():java.util.Set org.apache.doris.connector.spi.write.ConnectorWritePlanProvider#getSyntheticWriteColumns(org.apache.doris.connector.spi.ConnectorSession,org.apache.doris.connector.spi.handle.ConnectorTableHandle):java.util.List org.apache.doris.connector.spi.write.ConnectorWritePlanProvider#getWriteColumns(org.apache.doris.connector.spi.ConnectorSession,org.apache.doris.connector.spi.handle.ConnectorTableHandle,java.util.Optional):java.util.Optional +org.apache.doris.connector.spi.write.ConnectorWritePlanProvider#getWriteDistribution(org.apache.doris.connector.spi.ConnectorSession,org.apache.doris.connector.spi.handle.ConnectorTableHandle):org.apache.doris.connector.spi.write.ConnectorWriteDistribution org.apache.doris.connector.spi.write.ConnectorWritePlanProvider#getWriteMetadataIdentity(org.apache.doris.connector.spi.ConnectorSession,org.apache.doris.connector.spi.handle.ConnectorTableHandle):java.lang.String org.apache.doris.connector.spi.write.ConnectorWritePlanProvider#getWritePartitioning(org.apache.doris.connector.spi.ConnectorSession,org.apache.doris.connector.spi.handle.ConnectorTableHandle):org.apache.doris.connector.spi.write.ConnectorWritePartitionSpec org.apache.doris.connector.spi.write.ConnectorWritePlanProvider#getWriteSortColumns(org.apache.doris.connector.spi.ConnectorSession,org.apache.doris.connector.spi.handle.ConnectorTableHandle):java.util.List @@ -176,6 +208,7 @@ org.apache.doris.connector.spi.write.ConnectorWritePlanProvider#requiresPartitio org.apache.doris.connector.spi.write.ConnectorWritePlanProvider#requiresPartitionLocalSort():boolean org.apache.doris.connector.spi.write.ConnectorWritePlanProvider#supportedOperations():java.util.Set org.apache.doris.connector.spi.write.ConnectorWritePlanProvider#supportsWriteBranch():boolean +org.apache.doris.connector.spi.write.ConnectorWritePlanProvider#validateRowLevelDml(org.apache.doris.connector.spi.ConnectorSession,org.apache.doris.connector.spi.handle.ConnectorTableHandle,org.apache.doris.connector.spi.write.ConnectorRowLevelDmlRequest):void org.apache.doris.extension.spi.Plugin#close():void org.apache.doris.extension.spi.Plugin#initialize(org.apache.doris.extension.spi.PluginContext):void org.apache.doris.extension.spi.PluginContext#getProperties():java.util.Map diff --git a/fe/fe-connector/pom.xml b/fe/fe-connector/pom.xml index a682069b38b0c8..003e3412d64f4f 100644 --- a/fe/fe-connector/pom.xml +++ b/fe/fe-connector/pom.xml @@ -55,7 +55,7 @@ under the License. of the latter two means bumping this property as well (and fe-extension-spi means bumping all five families). --> - 10.0 + 11.0 diff --git a/fe/fe-core/src/main/java/org/apache/doris/analysis/ExprToThriftVisitor.java b/fe/fe-core/src/main/java/org/apache/doris/analysis/ExprToThriftVisitor.java index 6cd954bc8a6d4c..7a7401982f3638 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/analysis/ExprToThriftVisitor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/analysis/ExprToThriftVisitor.java @@ -483,7 +483,9 @@ public Void visitFunctionCallExpr(FunctionCallExpr expr, TExprNode msg) { msg.node_type = TExprNodeType.FUNCTION_CALL; } - if (ConnectContext.get() != null) { + if (expr instanceof ShortCircuitFunctionCallExpr) { + msg.setShortCircuitEvaluation(true); + } else if (ConnectContext.get() != null) { msg.setShortCircuitEvaluation(ConnectContext.get().getSessionVariable().isShortCircuitEvaluation()); } return null; diff --git a/fe/fe-core/src/main/java/org/apache/doris/datasource/plugin/PluginDrivenExternalTable.java b/fe/fe-core/src/main/java/org/apache/doris/datasource/plugin/PluginDrivenExternalTable.java index dd0bc658e86330..abbdccf8466fc6 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/datasource/plugin/PluginDrivenExternalTable.java +++ b/fe/fe-core/src/main/java/org/apache/doris/datasource/plugin/PluginDrivenExternalTable.java @@ -39,6 +39,10 @@ import org.apache.doris.connector.spi.handle.WriteOperation; import org.apache.doris.connector.spi.mvcc.ConnectorMvccSnapshot; import org.apache.doris.connector.spi.pushdown.ConnectorExpression; +import org.apache.doris.connector.spi.write.ConnectorChangelogMode; +import org.apache.doris.connector.spi.write.ConnectorRowChangeStyle; +import org.apache.doris.connector.spi.write.ConnectorRowLevelDmlRequest; +import org.apache.doris.connector.spi.write.ConnectorWriteDistribution; import org.apache.doris.connector.spi.write.ConnectorWritePlanProvider; import org.apache.doris.datasource.ExternalCatalog; import org.apache.doris.datasource.ExternalDatabase; @@ -77,6 +81,7 @@ import java.util.Map.Entry; import java.util.Optional; import java.util.Set; +import java.util.function.Supplier; import java.util.stream.Collectors; /** @@ -175,15 +180,9 @@ public boolean supportsParallelWrite() { ConnectorWritePlanProvider provider = writePlanProvider(); // requiresParallelWrite is byte-inert for a heterogeneous gateway (hive and iceberg both true), so the // connector-level answer needs no per-handle resolution here. - return provider != null && provider.requiresParallelWrite(); + return provider != null && withPluginContextClassLoader(provider, provider::requiresParallelWrite); } - /** - * Resolves this table's connector handle for a per-handle write-capability probe, or empty on any miss (a - * null connector, or an unresolvable handle). A heterogeneous gateway needs the handle to answer write - * capabilities per-table (its iceberg tables differ from its hive tables); a single-format connector ignores - * the handle (the per-handle overloads default to connector-level), so this is byte-identical for it. - */ /** * The CONNECTOR-LEVEL write plan provider, or null when this catalog's connector is absent or declares no * write support. Callers must have already checked that the catalog is plugin-driven. Used by the write @@ -193,12 +192,35 @@ public boolean supportsParallelWrite() { */ private ConnectorWritePlanProvider writePlanProvider() { Connector connector = ((PluginDrivenExternalCatalog) catalog).getConnector(); - return connector == null ? null : connector.getWritePlanProvider(); + return connector == null ? null + : withPluginContextClassLoader(connector, connector::getWritePlanProvider); } private Optional resolveWriteCapabilityHandle(Connector connector) { - ConnectorSession session = ((PluginDrivenExternalCatalog) catalog).buildConnectorSession(); - return resolveConnectorTableHandle(session, PluginDrivenMetadata.get(session, connector)); + return withPluginContextClassLoader(connector, () -> { + ConnectorSession session = ((PluginDrivenExternalCatalog) catalog).buildConnectorSession(); + return resolveConnectorTableHandle(session, PluginDrivenMetadata.get(session, connector)); + }); + } + + private ConnectorWritePlanProvider writePlanProvider( + Connector connector, ConnectorTableHandle handle) { + return withPluginContextClassLoader(connector, () -> connector.getWritePlanProvider(handle)); + } + + private Optional resolveWritePlanProvider(Connector connector) { + return resolveWriteCapabilityHandle(connector) + .map(handle -> writePlanProvider(connector, handle)); + } + + private static T withPluginContextClassLoader(Object plugin, Supplier callback) { + ClassLoader previous = Thread.currentThread().getContextClassLoader(); + try { + Thread.currentThread().setContextClassLoader(plugin.getClass().getClassLoader()); + return callback.get(); + } finally { + Thread.currentThread().setContextClassLoader(previous); + } } /** @@ -214,12 +236,105 @@ public Set connectorSupportedWriteOperations() { if (connector == null) { return EnumSet.noneOf(WriteOperation.class); } - return resolveWriteCapabilityHandle(connector) - .map(connector::getWritePlanProvider) - .map(ConnectorWritePlanProvider::supportedOperations) + return resolveWritePlanProvider(connector) + .map(provider -> withPluginContextClassLoader(provider, provider::supportedOperations)) .orElseGet(() -> EnumSet.noneOf(WriteOperation.class)); } + /** Returns the row-change representation declared for this table's write provider. */ + public ConnectorRowChangeStyle getConnectorRowChangeStyle() { + if (!(catalog instanceof PluginDrivenExternalCatalog)) { + return ConnectorRowChangeStyle.NONE; + } + Connector connector = ((PluginDrivenExternalCatalog) catalog).getConnector(); + if (connector == null) { + return ConnectorRowChangeStyle.NONE; + } + return resolveWritePlanProvider(connector) + .map(provider -> withPluginContextClassLoader(provider, provider::getRowChangeStyle)) + .orElse(ConnectorRowChangeStyle.NONE); + } + + /** Returns the operation-column encoding declared for this table's changelog writes. */ + public Optional getConnectorChangelogMode() { + if (!(catalog instanceof PluginDrivenExternalCatalog)) { + return Optional.empty(); + } + Connector connector = ((PluginDrivenExternalCatalog) catalog).getConnector(); + if (connector == null) { + return Optional.empty(); + } + return resolveWritePlanProvider(connector) + .flatMap(provider -> withPluginContextClassLoader(provider, provider::getChangelogMode)); + } + + /** Returns the primary-key columns used by this table's changelog row-level plan. */ + public List getConnectorRowLevelPrimaryKeyColumns() { + PluginDrivenExternalCatalog pluginCatalog = (PluginDrivenExternalCatalog) catalog; + Connector connector = pluginCatalog.getConnector(); + return withPluginContextClassLoader(connector, () -> { + ConnectorSession session = pluginCatalog.buildConnectorSession(); + ConnectorMetadata metadata = PluginDrivenMetadata.get(session, connector); + ConnectorTableHandle handle = resolveConnectorTableHandle(session, metadata) + .orElseThrow(() -> new DorisConnectorException( + "Cannot resolve row-level DML target " + getName())); + ConnectorWritePlanProvider provider = writePlanProvider(connector, handle); + return withPluginContextClassLoader(provider, + () -> provider.getRowLevelPrimaryKeyColumns(session, handle)); + }); + } + + /** Runs the engine-neutral mode check and connector-specific row-level validation. */ + public void validateConnectorRowLevelDml(ConnectorRowLevelDmlRequest request) { + PluginDrivenExternalCatalog pluginCatalog = (PluginDrivenExternalCatalog) catalog; + Connector connector = pluginCatalog.getConnector(); + withPluginContextClassLoader(connector, () -> { + ConnectorSession session = pluginCatalog.buildConnectorSession(); + ConnectorMetadata metadata = PluginDrivenMetadata.get(session, connector); + ConnectorTableHandle handle = resolveConnectorTableHandle(session, metadata) + .orElseThrow(() -> new DorisConnectorException( + "Cannot resolve row-level DML target " + getName())); + metadata.validateRowLevelDmlMode(session, handle, request.getOperation()); + ConnectorWritePlanProvider provider = writePlanProvider(connector, handle); + withPluginContextClassLoader(provider, () -> { + provider.validateRowLevelDml(session, handle, request); + return null; + }); + return null; + }); + } + + /** Returns connector-declared synthetic columns excluded from row-level write constraints. */ + public Set getConnectorRowLevelWriteConstraintExcludedColumns() { + if (!(catalog instanceof PluginDrivenExternalCatalog)) { + return Collections.emptySet(); + } + Connector connector = ((PluginDrivenExternalCatalog) catalog).getConnector(); + if (connector == null) { + return Collections.emptySet(); + } + return resolveWritePlanProvider(connector) + .map(provider -> withPluginContextClassLoader(provider, + provider::getRowLevelWriteConstraintExcludedColumns)) + .orElseGet(Collections::emptySet); + } + + /** Returns the connector-owned transaction-label prefix for one row-level operation. */ + public String getConnectorRowLevelDmlLabelPrefix(WriteOperation operation) { + if (!(catalog instanceof PluginDrivenExternalCatalog)) { + throw new DorisConnectorException("Row-level DML requires a plugin-driven catalog"); + } + Connector connector = ((PluginDrivenExternalCatalog) catalog).getConnector(); + if (connector == null) { + throw new DorisConnectorException("Connector is unavailable for row-level DML"); + } + return resolveWritePlanProvider(connector) + .map(provider -> withPluginContextClassLoader(provider, + () -> provider.getRowLevelDmlLabelPrefix(operation))) + .orElseThrow(() -> new DorisConnectorException( + "Cannot resolve the connector write provider for row-level DML")); + } + /** * Whether the connector admits branch writes for THIS table, resolved per-handle (iceberg supports * write-to-branch, hive does not). Degrades to false on any miss. @@ -232,9 +347,8 @@ public boolean connectorSupportsWriteBranch() { if (connector == null) { return false; } - return resolveWriteCapabilityHandle(connector) - .map(connector::getWritePlanProvider) - .map(ConnectorWritePlanProvider::supportsWriteBranch) + return resolveWritePlanProvider(connector) + .map(provider -> withPluginContextClassLoader(provider, provider::supportsWriteBranch)) .orElse(false); } @@ -397,7 +511,8 @@ public boolean requirePartitionLocalSortOnWrite() { return false; } ConnectorWritePlanProvider provider = writePlanProvider(); - return provider != null && provider.requiresPartitionLocalSort(); + return provider != null + && withPluginContextClassLoader(provider, provider::requiresPartitionLocalSort); } /** @@ -417,12 +532,36 @@ public boolean requirePartitionHashOnWrite() { return false; } // Per-table: hive requires partition-hash writes but iceberg does not, so resolve the handle. - return resolveWriteCapabilityHandle(connector) - .map(connector::getWritePlanProvider) - .map(ConnectorWritePlanProvider::requiresPartitionHashWrite) + return resolveWritePlanProvider(connector) + .map(provider -> withPluginContextClassLoader(provider, + provider::requiresPartitionHashWrite)) .orElse(false); } + /** Returns this table's connector-owned write distribution, or empty for generic planning. */ + public Optional getConnectorWriteDistribution() { + if (!(catalog instanceof PluginDrivenExternalCatalog)) { + return Optional.empty(); + } + PluginDrivenExternalCatalog pluginCatalog = (PluginDrivenExternalCatalog) catalog; + Connector connector = pluginCatalog.getConnector(); + if (connector == null) { + return Optional.empty(); + } + ConnectorSession session = pluginCatalog.buildConnectorSession(); + ConnectorMetadata metadata = PluginDrivenMetadata.get(session, connector); + Optional handle = resolveConnectorTableHandle(session, metadata); + if (!handle.isPresent()) { + return Optional.empty(); + } + ConnectorWritePlanProvider provider = writePlanProvider(connector, handle.get()); + if (provider == null) { + return Optional.empty(); + } + return Optional.ofNullable(withPluginContextClassLoader(provider, + () -> provider.getWriteDistribution(session, handle.get()))); + } + /** * Returns whether the underlying connector maps write data columns positionally against the full * table schema (e.g. MaxCompute), requiring the sink to project rows to full-schema order with @@ -434,7 +573,8 @@ public boolean requiresFullSchemaWriteOrder() { return false; } ConnectorWritePlanProvider provider = writePlanProvider(); - return provider != null && provider.requiresFullSchemaWriteOrder(); + return provider != null + && withPluginContextClassLoader(provider, provider::requiresFullSchemaWriteOrder); } /** @@ -454,9 +594,9 @@ public boolean materializeStaticPartitionValues() { } // Per-table: iceberg retains partition columns and hive derives the partition directory from the row // (both materialize the PARTITION literal); maxcompute refills from the static spec instead. - return resolveWriteCapabilityHandle(connector) - .map(connector::getWritePlanProvider) - .map(ConnectorWritePlanProvider::requiresMaterializeStaticPartitionValues) + return resolveWritePlanProvider(connector) + .map(provider -> withPluginContextClassLoader(provider, + provider::requiresMaterializeStaticPartitionValues)) .orElse(false); } @@ -825,11 +965,12 @@ private List fetchSyntheticWriteColumns() { if (!handleOpt.isPresent()) { return Collections.emptyList(); } - ConnectorWritePlanProvider writePlanProvider = connector.getWritePlanProvider(handleOpt.get()); + ConnectorWritePlanProvider writePlanProvider = writePlanProvider(connector, handleOpt.get()); if (writePlanProvider == null) { return Collections.emptyList(); } - return writePlanProvider.getSyntheticWriteColumns(session, handleOpt.get()); + return withPluginContextClassLoader(writePlanProvider, + () -> writePlanProvider.getSyntheticWriteColumns(session, handleOpt.get())); } /** @@ -859,13 +1000,11 @@ public Optional> resolveWriteColumns(Optional branchName) { if (!handle.isPresent()) { return Optional.empty(); } - ConnectorWritePlanProvider provider = connector.getWritePlanProvider(handle.get()); + ConnectorWritePlanProvider provider = writePlanProvider(connector, handle.get()); if (provider == null) { return Optional.empty(); } - ClassLoader previous = Thread.currentThread().getContextClassLoader(); - try { - Thread.currentThread().setContextClassLoader(provider.getClass().getClassLoader()); + return withPluginContextClassLoader(provider, () -> { Optional> connectorColumns = provider.getWriteColumns(session, handle.get(), branchName); if (!connectorColumns.isPresent()) { @@ -877,9 +1016,7 @@ public Optional> resolveWriteColumns(Optional branchName) { ctx.getStatementContext().setConnectorWriteMetadataIdentity(getId(), identity); } return connectorColumns.map(ConnectorColumnConverter::convertColumns); - } finally { - Thread.currentThread().setContextClassLoader(previous); - } + }); } /** The raw connector-emitted table-property map (including FE-internal / render-hint keys). */ diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/analyzer/UnboundConnectorTableSink.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/analyzer/UnboundConnectorTableSink.java index f5e4be4182d204..70f0020b9a140f 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/analyzer/UnboundConnectorTableSink.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/analyzer/UnboundConnectorTableSink.java @@ -22,6 +22,7 @@ import org.apache.doris.nereids.trees.expressions.Expression; import org.apache.doris.nereids.trees.plans.Plan; import org.apache.doris.nereids.trees.plans.PlanType; +import org.apache.doris.nereids.trees.plans.commands.info.ConnectorChangelogRowChangeSpec; import org.apache.doris.nereids.trees.plans.commands.info.DMLCommandType; import org.apache.doris.nereids.trees.plans.visitor.PlanVisitor; @@ -31,6 +32,7 @@ import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.Optional; /** @@ -49,6 +51,7 @@ public class UnboundConnectorTableSink extends UnboundB // rewrite_data_files INSERT-SELECT (controls output file count). Defaults false; set true only by the // distributed rewrite coordinator. Always false for ordinary INSERT, so this is dormant pre-cutover. private final boolean rewrite; + private final Optional rowChangeSpec; public UnboundConnectorTableSink(List nameParts, List colNames, List hints, List partitions, CHILD_TYPE child) { @@ -81,7 +84,7 @@ public UnboundConnectorTableSink(List nameParts, CHILD_TYPE child, Map staticPartitionKeyValues) { this(nameParts, colNames, hints, partitions, dmlCommandType, - groupExpression, logicalProperties, child, staticPartitionKeyValues, false); + groupExpression, logicalProperties, child, staticPartitionKeyValues, false, Optional.empty()); } /** @@ -97,12 +100,34 @@ public UnboundConnectorTableSink(List nameParts, CHILD_TYPE child, Map staticPartitionKeyValues, boolean rewrite) { + this(nameParts, colNames, hints, partitions, dmlCommandType, groupExpression, + logicalProperties, child, staticPartitionKeyValues, rewrite, Optional.empty()); + } + + /** Creates an unbound connector changelog sink for row-level DML. */ + public UnboundConnectorTableSink(List nameParts, CHILD_TYPE child, + ConnectorChangelogRowChangeSpec rowChangeSpec) { + this(nameParts, ImmutableList.of(), ImmutableList.of(), ImmutableList.of(), + rowChangeSpec.getDmlCommandType(), Optional.empty(), Optional.empty(), child, + null, false, Optional.of(rowChangeSpec)); + } + + private UnboundConnectorTableSink(List nameParts, List colNames, + List hints, List partitions, + DMLCommandType dmlCommandType, + Optional groupExpression, + Optional logicalProperties, + CHILD_TYPE child, + Map staticPartitionKeyValues, + boolean rewrite, + Optional rowChangeSpec) { super(nameParts, PlanType.LOGICAL_UNBOUND_CONNECTOR_TABLE_SINK, ImmutableList.of(), groupExpression, logicalProperties, colNames, dmlCommandType, child, hints, partitions); this.staticPartitionKeyValues = staticPartitionKeyValues != null ? ImmutableMap.copyOf(staticPartitionKeyValues) : null; this.rewrite = rewrite; + this.rowChangeSpec = rowChangeSpec; } public Map getStaticPartitionKeyValues() { @@ -117,6 +142,28 @@ public boolean hasStaticPartition() { return staticPartitionKeyValues != null && !staticPartitionKeyValues.isEmpty(); } + public Optional getRowChangeSpec() { + return rowChangeSpec; + } + + @Override + public List getExpressions() { + return rowChangeSpec.isPresent() ? rowChangeSpec.get().getExpressions() : super.getExpressions(); + } + + @Override + public boolean equals(Object other) { + return other instanceof UnboundConnectorTableSink + && super.equals(other) + && Objects.equals(rowChangeSpec, + ((UnboundConnectorTableSink) other).rowChangeSpec); + } + + @Override + public int hashCode() { + return Objects.hash(super.hashCode(), rowChangeSpec); + } + @Override public R accept(PlanVisitor visitor, C context) { return visitor.visitUnboundConnectorTableSink(this, context); @@ -127,20 +174,22 @@ public Plan withChildren(List children) { Preconditions.checkArgument(children.size() == 1, "UnboundConnectorTableSink only accepts one child"); return new UnboundConnectorTableSink<>(nameParts, colNames, hints, partitions, - dmlCommandType, groupExpression, Optional.empty(), children.get(0), staticPartitionKeyValues, rewrite); + dmlCommandType, groupExpression, Optional.empty(), children.get(0), staticPartitionKeyValues, + rewrite, rowChangeSpec); } @Override public Plan withGroupExpression(Optional groupExpression) { return new UnboundConnectorTableSink<>(nameParts, colNames, hints, partitions, dmlCommandType, groupExpression, Optional.of(getLogicalProperties()), child(), - staticPartitionKeyValues, rewrite); + staticPartitionKeyValues, rewrite, rowChangeSpec); } @Override public Plan withGroupExprLogicalPropChildren(Optional groupExpression, Optional logicalProperties, List children) { return new UnboundConnectorTableSink<>(nameParts, colNames, hints, partitions, - dmlCommandType, groupExpression, logicalProperties, children.get(0), staticPartitionKeyValues, rewrite); + dmlCommandType, groupExpression, logicalProperties, children.get(0), staticPartitionKeyValues, + rewrite, rowChangeSpec); } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/ExpressionTranslator.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/ExpressionTranslator.java index 162bece0228660..1b3d0f3990492f 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/ExpressionTranslator.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/ExpressionTranslator.java @@ -36,6 +36,7 @@ import org.apache.doris.analysis.MatchPredicate; import org.apache.doris.analysis.OrderByElement; import org.apache.doris.analysis.SearchPredicate; +import org.apache.doris.analysis.ShortCircuitFunctionCallExpr; import org.apache.doris.analysis.SlotDescriptor; import org.apache.doris.analysis.SlotRef; import org.apache.doris.analysis.TryCastExpr; @@ -85,6 +86,7 @@ import org.apache.doris.nereids.trees.expressions.functions.AlwaysNotNullable; import org.apache.doris.nereids.trees.expressions.functions.AlwaysNullable; import org.apache.doris.nereids.trees.expressions.functions.PropagateNullLiteral; +import org.apache.doris.nereids.trees.expressions.functions.RequiresShortCircuitEvaluation; import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction; import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateParam; import org.apache.doris.nereids.trees.expressions.functions.agg.Count; @@ -744,7 +746,10 @@ public Expr visitScalarFunction(ScalarFunction function, PlanTranslatorContext c "", Function.BinaryType.BUILTIN, true, true, nullableMode); // create catalog FunctionCallExpr without analyze again - return new FunctionCallExpr(catalogFunction, new FunctionParams(false, arguments), function.nullable()); + FunctionParams functionParams = new FunctionParams(false, arguments); + return function instanceof RequiresShortCircuitEvaluation + ? new ShortCircuitFunctionCallExpr(catalogFunction, functionParams, function.nullable()) + : new FunctionCallExpr(catalogFunction, functionParams, function.nullable()); } @Override diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java index 9bfe285a13e645..57e8ecad9c5e86 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/glue/translator/PhysicalPlanTranslator.java @@ -66,6 +66,8 @@ import org.apache.doris.nereids.properties.DistributionSpecAllSingleton; import org.apache.doris.nereids.properties.DistributionSpecAny; import org.apache.doris.nereids.properties.DistributionSpecExecutionAny; +import org.apache.doris.nereids.properties.DistributionSpecExternalTableSinkHashPartitioned; +import org.apache.doris.nereids.properties.DistributionSpecExternalTableSinkUnPartitioned; import org.apache.doris.nereids.properties.DistributionSpecGather; import org.apache.doris.nereids.properties.DistributionSpecHash; import org.apache.doris.nereids.properties.DistributionSpecHiveTableSinkHashPartitioned; @@ -218,6 +220,7 @@ import org.apache.doris.statistics.StatisticConstants; import org.apache.doris.tablefunction.TableValuedFunctionIf; import org.apache.doris.thrift.TBinlogScanType; +import org.apache.doris.thrift.TExternalTableSinkWriterAssignment; import org.apache.doris.thrift.TPartitionType; import org.apache.doris.thrift.TPushAggOp; import org.apache.doris.thrift.TResultSinkType; @@ -709,13 +712,14 @@ public PlanFragment visitPhysicalConnectorTableSink( "Table not found: " + targetTable.getRemoteDbName() + "." + targetTable.getRemoteName() + " in catalog " + catalog.getName())); - // Resolve the provider once: it both admits INSERT and plans the sink (see the row-level DML arm). + // Resolve the provider once: it both admits this write operation and plans the sink. ConnectorWritePlanProvider writePlanProvider = connector.getWritePlanProvider(providerTableHandle); + WriteOperation writeOperation = connectorWriteOperation(connectorTableSink); if (writePlanProvider == null - || !writePlanProvider.supportedOperations().contains(WriteOperation.INSERT)) { + || !writePlanProvider.supportedOperations().contains(writeOperation)) { throw new AnalysisException( "Connector '" + catalog.getName() + "' (type: " + catalog.getType() - + ") does not support INSERT operations"); + + ") does not support " + writeOperation + " operations"); } // Preserve the generation captured from the exact remote table load that supplied the bound schema. @@ -729,12 +733,6 @@ public PlanFragment visitPhysicalConnectorTableSink( writePlanProvider.getWriteSortColumns(connSession, providerTableHandle, boundOutputColumns), connectorTableSink, context); - // A distributed rewrite_data_files INSERT-SELECT threads WriteOperation.REWRITE so the connector's - // planWrite enters its REWRITE arm (RewriteFiles semantics) instead of the plain-INSERT append; the - // rewrite marker rides on the sink (PhysicalConnectorTableSink.isRewrite), not on a ConnectContext or - // an instanceof Iceberg. Ordinary connector INSERTs keep WriteOperation.INSERT (byte-identical). - WriteOperation writeOperation = connectorTableSink.isRewrite() - ? WriteOperation.REWRITE : WriteOperation.INSERT; // The write list can omit explicit/static-partition columns, but schema-drift validation must // retain the complete generation captured by BindSink instead of comparing that subset. PluginDrivenTableSink providerSink = new PluginDrivenTableSink(targetTable, @@ -745,6 +743,24 @@ public PlanFragment visitPhysicalConnectorTableSink( return rootFragment; } + private WriteOperation connectorWriteOperation(PhysicalConnectorTableSink sink) { + if (sink.getDmlCommandType() == null) { + // Legacy connector sinks do not carry a DML command type. Preserve their existing + // INSERT/REWRITE admission behavior while row-level sinks pass an explicit type. + return sink.isRewrite() ? WriteOperation.REWRITE : WriteOperation.INSERT; + } + switch (sink.getDmlCommandType()) { + case DELETE: + return WriteOperation.DELETE; + case UPDATE: + return WriteOperation.UPDATE; + case MERGE: + return WriteOperation.MERGE; + default: + return sink.isRewrite() ? WriteOperation.REWRITE : WriteOperation.INSERT; + } + } + private static ConnectorColumn toWriteConnectorColumn(Column column) { // Use the shared recursive conversion so write validation receives nested field identities as well // as the root id; rebuilding only the root silently accepted drop-and-recreate nested fields. @@ -762,9 +778,11 @@ private TSortInfo buildConnectorWriteSortInfo(List sor List orderingExprs = Lists.newArrayList(); List isAscOrder = Lists.newArrayList(); List nullsFirst = Lists.newArrayList(); + int outputOffset = connectorTableSink.hasRowOperationColumn() ? 1 : 0; for (ConnectorWriteSortColumn sortColumn : sortColumns) { orderingExprs.add(context.findSlotRef( - connectorTableSink.getOutput().get(sortColumn.getColumnIndex()).getExprId())); + connectorTableSink.getOutput().get( + sortColumn.getColumnIndex() + outputOffset).getExprId())); isAscOrder.add(sortColumn.isAsc()); nullsFirst.add(sortColumn.isNullsFirst()); } @@ -3749,6 +3767,35 @@ private DataPartition toDataPartition(DistributionSpec distributionSpec/* target return new DataPartition(partitionType, partitionExprs); } else if (distributionSpec instanceof DistributionSpecOlapTableSinkHashPartitioned) { return DataPartition.TABLET_ID; + } else if (distributionSpec instanceof DistributionSpecExternalTableSinkHashPartitioned) { + DistributionSpecExternalTableSinkHashPartitioned externalSpec + = (DistributionSpecExternalTableSinkHashPartitioned) distributionSpec; + List partitionExprs = Lists.newArrayList(); + for (ExprId partitionExprId : externalSpec.getOutputColumnExprIds()) { + if (childOutputIds.contains(partitionExprId)) { + partitionExprs.add(context.findSlotRef(partitionExprId)); + } + } + Preconditions.checkState(partitionExprs.size() + == externalSpec.getOutputColumnExprIds().size(), + "External sink route expressions must be present in child output"); + TExternalTableSinkWriterAssignment writerAssignment; + switch (externalSpec.getWriterAssignment()) { + case IDENTITY: + writerAssignment = TExternalTableSinkWriterAssignment.IDENTITY; + break; + case SKEWED: + writerAssignment = TExternalTableSinkWriterAssignment.SKEWED; + break; + default: + throw new IllegalStateException("Unsupported external sink writer assignment: " + + externalSpec.getWriterAssignment()); + } + return new DataPartition(TPartitionType.EXTERNAL_TABLE_SINK_HASH_PARTITIONED, + partitionExprs, externalSpec.getPartitionFunction(), + externalSpec.getPartitionFunctionOptions(), writerAssignment); + } else if (distributionSpec instanceof DistributionSpecExternalTableSinkUnPartitioned) { + return new DataPartition(TPartitionType.EXTERNAL_TABLE_SINK_UNPARTITIONED); } else if (distributionSpec instanceof DistributionSpecHiveTableSinkHashPartitioned) { DistributionSpecHiveTableSinkHashPartitioned partitionSpecHash = (DistributionSpecHiveTableSinkHashPartitioned) distributionSpec; diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/processor/post/CommonSubExpressionCollector.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/processor/post/CommonSubExpressionCollector.java index cccf9dbba072b9..98092a21560d63 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/processor/post/CommonSubExpressionCollector.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/processor/post/CommonSubExpressionCollector.java @@ -22,6 +22,7 @@ import org.apache.doris.nereids.trees.expressions.Expression; import org.apache.doris.nereids.trees.expressions.SessionVarGuardExpr; import org.apache.doris.nereids.trees.expressions.WhenClause; +import org.apache.doris.nereids.trees.expressions.functions.RequiresShortCircuitEvaluation; import org.apache.doris.nereids.trees.expressions.functions.scalar.Lambda; import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor; @@ -44,6 +45,9 @@ public int collect(Expression expr) { @Override public Integer visit(Expression expr, Boolean inLambda) { + if (expr instanceof RequiresShortCircuitEvaluation) { + return 0; + } return processExpressionWithChildren(expr.children(), expr, inLambda); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/DistributionSpecExternalTableSinkHashPartitioned.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/DistributionSpecExternalTableSinkHashPartitioned.java new file mode 100644 index 00000000000000..76a9d8523e692a --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/DistributionSpecExternalTableSinkHashPartitioned.java @@ -0,0 +1,91 @@ +// 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. + +package org.apache.doris.nereids.properties; + +import org.apache.doris.connector.spi.write.ConnectorWriteDistribution.WriterAssignment; +import org.apache.doris.nereids.trees.expressions.ExprId; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; + +import java.util.List; +import java.util.Map; +import java.util.Objects; + +/** Opaque connector partition-function request for an external table sink. */ +public final class DistributionSpecExternalTableSinkHashPartitioned extends DistributionSpec { + + public static final int MIN_BE_EXEC_VERSION = 13; + + private final ImmutableList outputColumnExprIds; + private final String partitionFunction; + private final ImmutableMap partitionFunctionOptions; + private final WriterAssignment writerAssignment; + + public DistributionSpecExternalTableSinkHashPartitioned(List outputColumnExprIds, + String partitionFunction, Map partitionFunctionOptions, + WriterAssignment writerAssignment) { + this.outputColumnExprIds = ImmutableList.copyOf(outputColumnExprIds); + this.partitionFunction = Objects.requireNonNull(partitionFunction); + this.partitionFunctionOptions = ImmutableMap.copyOf(partitionFunctionOptions); + this.writerAssignment = Objects.requireNonNull(writerAssignment); + } + + public List getOutputColumnExprIds() { + return outputColumnExprIds; + } + + public String getPartitionFunction() { + return partitionFunction; + } + + public Map getPartitionFunctionOptions() { + return partitionFunctionOptions; + } + + public WriterAssignment getWriterAssignment() { + return writerAssignment; + } + + @Override + public boolean satisfy(DistributionSpec required) { + return required instanceof DistributionSpecAny || equals(required); + } + + @Override + public boolean equals(Object other) { + if (this == other) { + return true; + } + if (!(other instanceof DistributionSpecExternalTableSinkHashPartitioned)) { + return false; + } + DistributionSpecExternalTableSinkHashPartitioned that + = (DistributionSpecExternalTableSinkHashPartitioned) other; + return outputColumnExprIds.equals(that.outputColumnExprIds) + && partitionFunction.equals(that.partitionFunction) + && partitionFunctionOptions.equals(that.partitionFunctionOptions) + && writerAssignment == that.writerAssignment; + } + + @Override + public int hashCode() { + return Objects.hash(outputColumnExprIds, partitionFunction, + partitionFunctionOptions, writerAssignment); + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/DistributionSpecExternalTableSinkUnPartitioned.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/DistributionSpecExternalTableSinkUnPartitioned.java new file mode 100644 index 00000000000000..5e17bf5b3a4a69 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/DistributionSpecExternalTableSinkUnPartitioned.java @@ -0,0 +1,33 @@ +// 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. + +package org.apache.doris.nereids.properties; + +/** Adaptive writer distribution for an external table without an ownership key. */ +public final class DistributionSpecExternalTableSinkUnPartitioned extends DistributionSpec { + + public static final DistributionSpecExternalTableSinkUnPartitioned INSTANCE + = new DistributionSpecExternalTableSinkUnPartitioned(); + + private DistributionSpecExternalTableSinkUnPartitioned() { + } + + @Override + public boolean satisfy(DistributionSpec required) { + return required instanceof DistributionSpecExternalTableSinkUnPartitioned; + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/PhysicalProperties.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/PhysicalProperties.java index c28d6ac3cb4d47..c24f0b5712044f 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/PhysicalProperties.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/PhysicalProperties.java @@ -52,6 +52,9 @@ public class PhysicalProperties { public static PhysicalProperties SINK_RANDOM_PARTITIONED = new PhysicalProperties(DistributionSpecHiveTableSinkUnPartitioned.INSTANCE); + public static PhysicalProperties EXTERNAL_TABLE_SINK_UNPARTITIONED + = new PhysicalProperties(DistributionSpecExternalTableSinkUnPartitioned.INSTANCE); + // gather then broadcast to all BE with exact one instance public static PhysicalProperties ALL_SINGLETON = new PhysicalProperties(DistributionSpecAllSingleton.INSTANCE); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/RequestPropertyDeriver.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/RequestPropertyDeriver.java index f46d5ff5206016..19043470128781 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/RequestPropertyDeriver.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/RequestPropertyDeriver.java @@ -185,8 +185,11 @@ public Void visitPhysicalConnectorTableSink( // Connector does not support parallel write (e.g., JDBC, ES). // Always gather to a single writer for transactional safety. addRequestPropertyToChildren(PhysicalProperties.GATHER); - } else if (connectContext != null + } else if (PhysicalProperties.SINK_RANDOM_PARTITIONED.equals(requiredProps) + && connectContext != null && !connectContext.getSessionVariable().isEnableStrictConsistencyDml()) { + // Strict-consistency mode may relax only the generic random parallel-write preference. + // Connector routing, partition hashing, and local ordering are writer correctness contracts. addRequestPropertyToChildren(PhysicalProperties.ANY); } else { addRequestPropertyToChildren(requiredProps); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/BindSink.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/BindSink.java index c00dca0813f2c4..526d936375b483 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/BindSink.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/BindSink.java @@ -38,6 +38,7 @@ import org.apache.doris.connector.spi.ConnectorSession; import org.apache.doris.connector.spi.DorisConnectorException; import org.apache.doris.connector.spi.handle.ConnectorTableHandle; +import org.apache.doris.connector.spi.write.ConnectorChangelogMode; import org.apache.doris.datasource.ExternalDatabase; import org.apache.doris.datasource.doris.RemoteDorisExternalTable; import org.apache.doris.datasource.plugin.PluginDrivenExternalCatalog; @@ -747,7 +748,7 @@ private static List sinkTargetFullSchema(TableIf table) { * stay in the connector (iceberg). A connector {@link DorisConnectorException} is surfaced as the * analysis-time {@link AnalysisException} the legacy native path threw, preserving the user-facing message * and the exception type. The literal-value check is connector-agnostic and stays here, where the Nereids - * expression is available. Plumbing mirrors {@code IcebergRowLevelDmlTransform.checkPluginMode}. + * expression is available. Plumbing mirrors {@code PositionDeleteRowLevelDmlTransform.checkPluginMode}. */ private void checkConnectorStaticPartitions(PluginDrivenExternalTable table, Map staticPartitions, Set staticPartitionColNames) { @@ -895,6 +896,26 @@ private Plan bindConnectorTableSink(MatchingContext targetWriteSchema = resolvedTargetSchema.stream() .filter(column -> isConnectorSinkWriteColumn(column, sink.isRewrite())) .collect(ImmutableList.toImmutableList()); + if (sink.getRowChangeSpec().isPresent()) { + ConnectorChangelogMode changelogMode = table.getConnectorChangelogMode() + .orElseThrow(() -> new AnalysisException( + "Connector changelog write mode is not configured for table " + table.getName())); + child = ConnectorChangelogPlanBuilder.build(targetWriteSchema, + table.getConnectorRowLevelPrimaryKeyColumns(), changelogMode, + sink.getRowChangeSpec().get(), child, ctx.cascadesContext); + List outputExpressions = child.getOutput().stream() + .map(NamedExpression.class::cast) + .collect(ImmutableList.toImmutableList()); + if (outputExpressions.size() != targetWriteSchema.size() + 1) { + throw new AnalysisException("Connector changelog sink must produce an operation column and " + + targetWriteSchema.size() + " table columns, but got " + outputExpressions.size()); + } + return new LogicalConnectorTableSink<>(database, table, targetWriteSchema, + targetMetadata.getPartitionColumns(), targetMetadata.getWriteMetadataIdentity(), + targetWriteSchema, outputExpressions, sink.getDMLCommandType(), false, + true, + Optional.empty(), Optional.empty(), child); + } if (sink.isRewrite()) { List rewriteOutputs = selectConnectorRewriteOutputs( targetWriteSchema, child.getOutput()); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/ConnectorChangelogPlanBuilder.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/ConnectorChangelogPlanBuilder.java new file mode 100644 index 00000000000000..4221cd28bfa4f2 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/analysis/ConnectorChangelogPlanBuilder.java @@ -0,0 +1,563 @@ +// 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. + +package org.apache.doris.nereids.rules.analysis; + +import org.apache.doris.catalog.Column; +import org.apache.doris.connector.spi.write.ConnectorChangelogMode; +import org.apache.doris.nereids.CascadesContext; +import org.apache.doris.nereids.analyzer.Scope; +import org.apache.doris.nereids.analyzer.UnboundAlias; +import org.apache.doris.nereids.analyzer.UnboundSlot; +import org.apache.doris.nereids.exceptions.AnalysisException; +import org.apache.doris.nereids.trees.expressions.Alias; +import org.apache.doris.nereids.trees.expressions.Cast; +import org.apache.doris.nereids.trees.expressions.EqualTo; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.LessThanEqual; +import org.apache.doris.nereids.trees.expressions.NamedExpression; +import org.apache.doris.nereids.trees.expressions.Not; +import org.apache.doris.nereids.trees.expressions.Slot; +import org.apache.doris.nereids.trees.expressions.WindowExpression; +import org.apache.doris.nereids.trees.expressions.functions.agg.AnyValue; +import org.apache.doris.nereids.trees.expressions.functions.agg.Count; +import org.apache.doris.nereids.trees.expressions.functions.scalar.AssertTrue; +import org.apache.doris.nereids.trees.expressions.literal.BigIntLiteral; +import org.apache.doris.nereids.trees.expressions.literal.IntegerLiteral; +import org.apache.doris.nereids.trees.expressions.literal.NullLiteral; +import org.apache.doris.nereids.trees.expressions.literal.TinyIntLiteral; +import org.apache.doris.nereids.trees.expressions.literal.VarcharLiteral; +import org.apache.doris.nereids.trees.plans.commands.ConnectorWriteSchemaUtils; +import org.apache.doris.nereids.trees.plans.commands.info.ConnectorChangelogRowChangeSpec; +import org.apache.doris.nereids.trees.plans.commands.merge.MergeMatchedClause; +import org.apache.doris.nereids.trees.plans.commands.merge.MergeNotMatchedClause; +import org.apache.doris.nereids.trees.plans.commands.merge.MergeUtils; +import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate; +import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; +import org.apache.doris.nereids.trees.plans.logical.LogicalJoin; +import org.apache.doris.nereids.trees.plans.logical.LogicalPlan; +import org.apache.doris.nereids.trees.plans.logical.LogicalProject; +import org.apache.doris.nereids.trees.plans.logical.LogicalWindow; +import org.apache.doris.nereids.types.BigIntType; +import org.apache.doris.nereids.types.DataType; +import org.apache.doris.nereids.types.IntegerType; +import org.apache.doris.nereids.util.ExpressionUtils; +import org.apache.doris.nereids.util.TypeCoercionUtils; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableSet; +import com.google.common.collect.Lists; +import com.google.common.collect.Maps; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.TreeSet; + +/** Builds the operation-column plus full-row projection used by changelog-oriented connectors. */ +public final class ConnectorChangelogPlanBuilder { + private static final String BRANCH_LABEL = "__DORIS_CHANGELOG_BRANCH__"; + + private ConnectorChangelogPlanBuilder() { + } + + /** Builds a changelog plan for the requested connector row-level operation. */ + public static LogicalPlan build(List schema, List primaryKeys, + ConnectorChangelogMode mode, ConnectorChangelogRowChangeSpec spec, + LogicalPlan child, CascadesContext context) { + if (spec instanceof ConnectorChangelogRowChangeSpec.Update) { + return buildUpdate(schema, mode, (ConnectorChangelogRowChangeSpec.Update) spec, + child, context); + } + if (spec instanceof ConnectorChangelogRowChangeSpec.Delete) { + return buildDelete(schema, primaryKeys, mode, + (ConnectorChangelogRowChangeSpec.Delete) spec, child, context); + } + if (spec instanceof ConnectorChangelogRowChangeSpec.Merge) { + return new MergeBuilder(schema, primaryKeys, mode, + (ConnectorChangelogRowChangeSpec.Merge) spec, child, context).build(); + } + throw new AnalysisException("Unsupported connector changelog specification: " + + spec.getClass().getSimpleName()); + } + + private static LogicalPlan buildUpdate(List schema, ConnectorChangelogMode mode, + ConnectorChangelogRowChangeSpec.Update update, LogicalPlan child, + CascadesContext context) { + Map changes = Maps.newTreeMap(String.CASE_INSENSITIVE_ORDER); + for (EqualTo assignment : update.getAssignments()) { + List parts = ((UnboundSlot) assignment.left()).getNameParts(); + String name = parts.get(parts.size() - 1); + if (changes.put(name, assignment.right()) != null) { + throw new AnalysisException("Duplicate column name in connector UPDATE: " + name); + } + } + ExpressionAnalyzer analyzer = analyzer(child, context); + List projects = new ArrayList<>(); + projects.add(operation(mode.getOperationColumnName(), mode.getUpdateValue())); + for (Column column : schema) { + Expression value = changes.remove(column.getName()); + if (value == null) { + value = targetSlot(update.getTargetNameInPlan(), column.getName()); + } + projects.add(bindColumn(analyzer, value, column)); + } + if (!changes.isEmpty()) { + throw new AnalysisException("Unknown column in connector UPDATE: " + + String.join(", ", changes.keySet())); + } + return new LogicalProject<>(projects, child); + } + + private static LogicalPlan buildDelete(List schema, List primaryKeys, + ConnectorChangelogMode mode, + ConnectorChangelogRowChangeSpec.Delete delete, LogicalPlan child, + CascadesContext context) { + ExpressionAnalyzer analyzer = analyzer(child, context); + List projects = new ArrayList<>(); + projects.add(operation(mode.getOperationColumnName(), mode.getDeleteValue())); + for (Column column : schema) { + projects.add(bindColumn(analyzer, + targetSlot(delete.getTargetNameInPlan(), column.getName()), column)); + } + LogicalProject project = new LogicalProject<>(projects, child); + if (!delete.shouldDeduplicateTargetRows()) { + return project; + } + if (primaryKeys.isEmpty()) { + throw new AnalysisException("Connector DELETE USING requires a primary-key table"); + } + Set keys = new TreeSet<>(String.CASE_INSENSITIVE_ORDER); + keys.addAll(primaryKeys); + List groupBy = new ArrayList<>(); + List outputs = new ArrayList<>(); + Slot operation = project.getOutput().get(0); + groupBy.add(operation); + outputs.add(operation); + for (int i = 0; i < schema.size(); i++) { + Column column = schema.get(i); + Slot value = project.getOutput().get(i + 1); + if (keys.remove(column.getName())) { + groupBy.add(value); + outputs.add(value); + } else { + outputs.add(new Alias(new AnyValue(value), column.getName())); + } + } + if (!keys.isEmpty()) { + throw new AnalysisException("Unknown connector primary-key column: " + + String.join(", ", keys)); + } + return new LogicalAggregate<>(groupBy, outputs, project); + } + + private static Alias operation(String columnName, byte value) { + return new Alias(new TinyIntLiteral(value), columnName); + } + + private static UnboundSlot targetSlot(List qualifier, String column) { + List parts = new ArrayList<>(qualifier); + parts.add(column); + return new UnboundSlot(parts); + } + + private static ExpressionAnalyzer analyzer(LogicalPlan plan, CascadesContext context) { + return new ExpressionAnalyzer(plan, new Scope(plan.getOutput()), context, true, false); + } + + private static Alias bindColumn(ExpressionAnalyzer analyzer, Expression expression, Column column) { + Expression value = analyzer.analyze(expression); + value = TypeCoercionUtils.castIfNotSameType(value, DataType.fromCatalogType(column.getType())); + return new Alias(value, column.getName()); + } + + private static final class MergeBuilder { + private final List schema; + private final List primaryKeys; + private final ConnectorChangelogMode mode; + private final ConnectorChangelogRowChangeSpec.Merge merge; + private final LogicalPlan child; + private final ExpressionAnalyzer analyzer; + + private MergeBuilder(List schema, List primaryKeys, + ConnectorChangelogMode mode, + ConnectorChangelogRowChangeSpec.Merge merge, LogicalPlan child, + CascadesContext context) { + this.schema = schema; + this.primaryKeys = primaryKeys; + this.mode = mode; + this.merge = merge; + this.child = child; + this.analyzer = analyzer(child, context); + } + + private LogicalPlan build() { + if (primaryKeys.isEmpty()) { + throw new AnalysisException("Connector MERGE requires a primary-key table"); + } + Alias branch = bindBranchLabel(); + Slot branchSlot = branch.toSlot(); + List branchOutputs = new ArrayList<>(child.getOutput()); + branchOutputs.add(branch); + LogicalPlan selected = new LogicalProject<>(branchOutputs, child); + selected = new LogicalFilter<>( + ImmutableSet.of(new Not(new org.apache.doris.nereids.trees.expressions.IsNull(branchSlot))), + selected); + List> branches = buildBranchProjections(); + if (!merge.getNotMatchedClauses().isEmpty()) { + validateNotMatchedPrimaryKeys(branches); + } + List output = new ArrayList<>(); + for (int column = 0; column <= schema.size(); column++) { + DataType type = column == 0 + ? org.apache.doris.nereids.types.TinyIntType.INSTANCE + : DataType.fromCatalogType(schema.get(column - 1).getType()); + String name = column == 0 + ? mode.getOperationColumnName() : schema.get(column - 1).getName(); + Expression value = new NullLiteral(type); + for (int index = branches.size() - 1; index >= 0; index--) { + Expression branchValue = TypeCoercionUtils.castIfNotSameType( + branches.get(index).get(column), type); + value = MergeUtils.selectBranch( + new EqualTo(branchSlot, new IntegerLiteral(index)), branchValue, value); + } + output.add(new Alias(value, name)); + } + int visibleOutputCount = output.size(); + if (!merge.getMatchedClauses().isEmpty()) { + for (String key : primaryKeys) { + output.add(new Alias(findTargetSlot(key), + "__DORIS_CHANGELOG_TARGET_KEY_" + key + "__")); + } + } + return addCardinalityChecks(new LogicalProject<>(output, selected), visibleOutputCount); + } + + private void validateNotMatchedPrimaryKeys(List> branches) { + Map targetKeys = Maps.newTreeMap(String.CASE_INSENSITIVE_ORDER); + for (String key : primaryKeys) { + targetKeys.put(key, findTargetSlot(key)); + } + Set targetSlots = child.getOutput().stream() + .filter(slot -> qualifierEndsWith(slot.getQualifier(), merge.getTargetNameInPlan())) + .collect(ImmutableSet.toImmutableSet()); + if (!(child instanceof LogicalJoin)) { + throw new AnalysisException("Connector MERGE input must be a logical join"); + } + Expression onClause = ((LogicalJoin) child).getOnClauseCondition() + .orElseThrow(() -> new AnalysisException("Connector MERGE requires an ON condition")); + Map sourceKeys = Maps.newTreeMap(String.CASE_INSENSITIVE_ORDER); + for (Expression conjunct : ExpressionUtils.extractConjunction(onClause)) { + if (!(conjunct instanceof EqualTo)) { + throw invalidNotMatchedKeyCondition(); + } + EqualTo equality = (EqualTo) conjunct; + String leftKey = targetPrimaryKeyName(equality.left(), targetKeys); + String rightKey = targetPrimaryKeyName(equality.right(), targetKeys); + if ((leftKey == null) == (rightKey == null)) { + throw invalidNotMatchedKeyCondition(); + } + String key = leftKey != null ? leftKey : rightKey; + Expression source = leftKey != null ? equality.right() : equality.left(); + if (source.getInputSlots().isEmpty() + || source.getInputSlots().stream().anyMatch(targetSlots::contains) + || source.containsNondeterministic() + || sourceKeys.put(key, source) != null) { + throw invalidNotMatchedKeyCondition(); + } + } + if (sourceKeys.size() != targetKeys.size()) { + throw invalidNotMatchedKeyCondition(); + } + int firstInsert = merge.getMatchedClauses().size(); + for (int branch = firstInsert; branch < branches.size(); branch++) { + for (Map.Entry sourceKey : sourceKeys.entrySet()) { + int column = schemaIndex(sourceKey.getKey()) + 1; + DataType type = DataType.fromCatalogType(schema.get(column - 1).getType()); + if (!TypeCoercionUtils.castIfNotSameType(branches.get(branch).get(column), type) + .equals(TypeCoercionUtils.castIfNotSameType(sourceKey.getValue(), type))) { + throw invalidNotMatchedKeyCondition(); + } + } + } + } + + private LogicalPlan addCardinalityChecks(LogicalProject rowChanges, int visibleOutputCount) { + List allOutputs = rowChanges.getOutput(); + List outputs = allOutputs.subList(0, visibleOutputCount); + Slot operation = outputs.get(0); + List insertedKeys = new ArrayList<>(); + for (String key : primaryKeys) { + insertedKeys.add(outputs.get(schemaIndex(key) + 1)); + } + List matchedKeys = new ArrayList<>( + allOutputs.subList(visibleOutputCount, allOutputs.size())); + Expression isInsert = new EqualTo(operation, new TinyIntLiteral(mode.getInsertValue())); + List checks = new ArrayList<>(); + if (!merge.getMatchedClauses().isEmpty()) { + checks.add(CardinalityCheck.matched(isInsert, matchedKeys)); + } + if (!merge.getNotMatchedClauses().isEmpty()) { + checks.add(CardinalityCheck.inserted(isInsert, insertedKeys)); + } + List markerOutputs = new ArrayList<>(allOutputs); + for (CardinalityCheck check : checks) { + markerOutputs.add(check.marker); + } + LogicalPlan plan = new LogicalProject<>(markerOutputs, rowChanges); + List counts = new ArrayList<>(); + for (CardinalityCheck check : checks) { + counts.add(check.count()); + } + plan = new LogicalWindow<>(new ArrayList<>(counts), plan); + ImmutableSet.Builder assertions = ImmutableSet.builder(); + for (int i = 0; i < checks.size(); i++) { + assertions.add(checks.get(i).assertion(counts.get(i))); + } + plan = new LogicalFilter<>(assertions.build(), plan); + return new LogicalProject<>(new ArrayList<>(outputs), plan); + } + + private int schemaIndex(String name) { + for (int i = 0; i < schema.size(); i++) { + if (schema.get(i).getName().equalsIgnoreCase(name)) { + return i; + } + } + throw new AnalysisException("Unable to resolve connector row-change column '" + name + "'"); + } + + private String targetPrimaryKeyName(Expression expression, Map targetKeys) { + Expression unwrapped = expression; + while (unwrapped instanceof Cast) { + if (((Cast) unwrapped).isExplicitType()) { + return null; + } + unwrapped = unwrapped.child(0); + } + if (!(unwrapped instanceof Slot)) { + return null; + } + Slot slot = (Slot) unwrapped; + for (Map.Entry key : targetKeys.entrySet()) { + if (slot.getExprId().equals(key.getValue().getExprId()) + && expression.getDataType().equals(key.getValue().getDataType())) { + return key.getKey(); + } + } + return null; + } + + private AnalysisException invalidNotMatchedKeyCondition() { + return new AnalysisException("Connector MERGE with NOT MATCHED INSERT requires ON to contain " + + "only equality predicates for every target primary-key column and each INSERT " + + "to use the corresponding deterministic source expression"); + } + + private Alias bindBranchLabel() { + Expression targetPresent = new Not(new org.apache.doris.nereids.trees.expressions.IsNull( + findTargetSlot(primaryKeys.get(0)))); + Expression matched = new NullLiteral(IntegerType.INSTANCE); + for (int i = merge.getMatchedClauses().size() - 1; i >= 0; i--) { + MergeMatchedClause clause = merge.getMatchedClauses().get(i); + if (i != merge.getMatchedClauses().size() - 1 && !clause.getCasePredicate().isPresent()) { + throw new AnalysisException("Only the last matched clause may omit its condition"); + } + Expression label = new IntegerLiteral(i); + matched = clause.getCasePredicate().isPresent() + ? MergeUtils.selectBranch(clause.getCasePredicate().get(), label, matched) : label; + } + Expression notMatched = new NullLiteral(IntegerType.INSTANCE); + for (int i = merge.getNotMatchedClauses().size() - 1; i >= 0; i--) { + MergeNotMatchedClause clause = merge.getNotMatchedClauses().get(i); + if (i != merge.getNotMatchedClauses().size() - 1 + && !clause.getCasePredicate().isPresent()) { + throw new AnalysisException("Only the last not matched clause may omit its condition"); + } + Expression label = new IntegerLiteral(i + merge.getMatchedClauses().size()); + notMatched = clause.getCasePredicate().isPresent() + ? MergeUtils.selectBranch(clause.getCasePredicate().get(), label, notMatched) : label; + } + return new Alias(analyzer.analyze( + MergeUtils.selectBranch(targetPresent, matched, notMatched)), BRANCH_LABEL); + } + + private List> buildBranchProjections() { + List> branches = new ArrayList<>(); + for (MergeMatchedClause clause : merge.getMatchedClauses()) { + branches.add(clause.isDelete() ? deleteProjection() : updateProjection(clause)); + } + for (MergeNotMatchedClause clause : merge.getNotMatchedClauses()) { + branches.add(insertProjection(clause)); + } + if (branches.isEmpty()) { + throw new AnalysisException("Connector MERGE requires at least one WHEN clause"); + } + for (List branch : branches) { + for (int i = 0; i < branch.size(); i++) { + branch.set(i, analyzer.analyze(branch.get(i))); + } + } + return branches; + } + + private List deleteProjection() { + List output = new ArrayList<>(); + output.add(new TinyIntLiteral(mode.getDeleteValue())); + for (Column column : schema) { + output.add(targetSlot(column.getName())); + } + return output; + } + + private List updateProjection(MergeMatchedClause clause) { + Map changes = Maps.newTreeMap(String.CASE_INSENSITIVE_ORDER); + for (EqualTo assignment : clause.getAssignments()) { + List parts = ((UnboundSlot) assignment.left()).getNameParts(); + String name = parts.get(parts.size() - 1); + if (changes.put(name, assignment.right()) != null) { + throw new AnalysisException("Duplicate column name in connector MERGE UPDATE: " + name); + } + } + List output = new ArrayList<>(); + output.add(new TinyIntLiteral(mode.getUpdateValue())); + for (Column column : schema) { + output.add(changes.containsKey(column.getName()) + ? changes.remove(column.getName()) : targetSlot(column.getName())); + } + if (!changes.isEmpty()) { + throw new AnalysisException("Unknown column in connector MERGE UPDATE: " + + String.join(", ", changes.keySet())); + } + return output; + } + + private List insertProjection(MergeNotMatchedClause clause) { + Map values = Maps.newTreeMap(String.CASE_INSENSITIVE_ORDER); + if (!clause.getColNames().isEmpty()) { + if (clause.getColNames().size() != clause.getRow().size()) { + throw new AnalysisException("Column count doesn't match value count"); + } + for (int i = 0; i < clause.getColNames().size(); i++) { + String column = clause.getColNames().get(i); + if (values.put(column, unwrap(clause.getRow().get(i))) != null) { + throw new AnalysisException("Duplicate column in connector MERGE INSERT: " + + column); + } + } + } else if (clause.getRow().size() != schema.size()) { + throw new AnalysisException("Column count doesn't match value count"); + } + List output = new ArrayList<>(); + output.add(new TinyIntLiteral(mode.getInsertValue())); + for (int i = 0; i < schema.size(); i++) { + Column column = schema.get(i); + Expression value = clause.getColNames().isEmpty() + ? unwrap(clause.getRow().get(i)) : values.remove(column.getName()); + if (value == null) { + value = ConnectorWriteSchemaUtils.resolveDefault(column); + } else { + value = ConnectorWriteSchemaUtils.resolveExplicitDefault(value, column); + } + output.add(value); + } + if (!values.isEmpty()) { + throw new AnalysisException("Unknown column in connector MERGE INSERT: " + + String.join(", ", values.keySet())); + } + return output; + } + + private Slot findTargetSlot(String column) { + List matches = child.getOutput().stream() + .filter(slot -> slot.getName().equalsIgnoreCase(column)) + .filter(slot -> qualifierEndsWith(slot.getQualifier(), merge.getTargetNameInPlan())) + .collect(java.util.stream.Collectors.toList()); + if (matches.size() != 1) { + throw new AnalysisException("Unable to resolve connector MERGE target column '" + + String.join(".", merge.getTargetNameInPlan()) + "." + column + "'"); + } + return matches.get(0); + } + + private Expression targetSlot(String column) { + List parts = Lists.newArrayList(merge.getTargetNameInPlan()); + parts.add(column); + return new UnboundSlot(parts); + } + + private static boolean qualifierEndsWith(List qualifier, List suffix) { + if (qualifier.size() < suffix.size()) { + return false; + } + int offset = qualifier.size() - suffix.size(); + for (int i = 0; i < suffix.size(); i++) { + if (!qualifier.get(offset + i).equalsIgnoreCase(suffix.get(i))) { + return false; + } + } + return true; + } + + private static Expression unwrap(NamedExpression expression) { + return expression instanceof Alias || expression instanceof UnboundAlias + ? expression.child(0) : expression; + } + + private static final class CardinalityCheck { + private final Alias marker; + private final List partitionKeys; + private final String countName; + private final String errorMessage; + + private CardinalityCheck(Alias marker, List partitionKeys, + String countName, String errorMessage) { + this.marker = marker; + this.partitionKeys = partitionKeys; + this.countName = countName; + this.errorMessage = errorMessage; + } + + private static CardinalityCheck matched(Expression isInsert, List partitionKeys) { + return new CardinalityCheck(new Alias(MergeUtils.selectBranch(isInsert, + new NullLiteral(BigIntType.INSTANCE), new BigIntLiteral(1)), + "__DORIS_CHANGELOG_MATCH_MARKER__"), partitionKeys, + "__DORIS_CHANGELOG_MATCH_COUNT__", + "Connector MERGE matched one target row with multiple source rows"); + } + + private static CardinalityCheck inserted(Expression isInsert, List partitionKeys) { + return new CardinalityCheck(new Alias(MergeUtils.selectBranch(isInsert, + new BigIntLiteral(1), new NullLiteral(BigIntType.INSTANCE)), + "__DORIS_CHANGELOG_INSERT_MARKER__"), partitionKeys, + "__DORIS_CHANGELOG_INSERT_COUNT__", + "Connector MERGE attempted to insert multiple rows with the same primary key"); + } + + private Alias count() { + return new Alias(new WindowExpression( + new Count(marker.toSlot()), partitionKeys, ImmutableList.of()), countName); + } + + private Expression assertion(Alias count) { + return new AssertTrue(new LessThanEqual(count.toSlot(), new BigIntLiteral(1)), + new VarcharLiteral(errorMessage)); + } + } + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/check/CheckCast.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/check/CheckCast.java index 23dffad6ff6c56..66a368c8315595 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/check/CheckCast.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/check/CheckCast.java @@ -29,6 +29,7 @@ import org.apache.doris.nereids.types.BitmapType; import org.apache.doris.nereids.types.BooleanType; import org.apache.doris.nereids.types.CharType; +import org.apache.doris.nereids.types.ConnectorComputeVariantType; import org.apache.doris.nereids.types.DataType; import org.apache.doris.nereids.types.DateTimeType; import org.apache.doris.nereids.types.DateTimeV2Type; @@ -384,6 +385,14 @@ public static boolean check(DataType originalType, DataType targetType, boolean */ public static boolean check(DataType originalType, DataType targetType, boolean isStrictMode, boolean looseAggState) { + if (originalType instanceof ConnectorComputeVariantType && targetType.isVariantType()) { + // The connector marker and ordinary Variant share the V2 runtime carrier. Allow the + // marker to cross the sink boundary without relaxing casts between stored Variant layouts. + return true; + } + if (targetType instanceof ConnectorComputeVariantType) { + return VariantType.isSupportedComputeV2CastSource(originalType); + } if (originalType.isVariantType() && (targetType instanceof PrimitiveType || targetType.isArrayType())) { // variant could cast to primitive types and array return true; diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/CaseWhenToCompoundPredicate.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/CaseWhenToCompoundPredicate.java index 51f4fce171af70..c90b3070c4d752 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/CaseWhenToCompoundPredicate.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/CaseWhenToCompoundPredicate.java @@ -27,6 +27,7 @@ import org.apache.doris.nereids.trees.expressions.NullSafeEqual; import org.apache.doris.nereids.trees.expressions.Or; import org.apache.doris.nereids.trees.expressions.WhenClause; +import org.apache.doris.nereids.trees.expressions.functions.RequiresShortCircuitEvaluation; import org.apache.doris.nereids.trees.expressions.functions.scalar.If; import org.apache.doris.nereids.trees.expressions.literal.BooleanLiteral; import org.apache.doris.nereids.trees.expressions.literal.NullLiteral; @@ -70,7 +71,12 @@ public List> buildRules() { } private boolean checkBooleanType(Expression expression) { - return expression.getDataType().isBooleanType(); + return expression.getDataType().isBooleanType() + && !requiresShortCircuitEvaluation(expression); + } + + private static boolean requiresShortCircuitEvaluation(Expression expression) { + return expression.anyMatch(node -> node instanceof RequiresShortCircuitEvaluation); } private Expression rewriteCaseWhen(CaseWhen caseWhen) { @@ -126,11 +132,15 @@ public List> buildRules() { @Override protected boolean needRewrite(Expression expression, boolean isInsideCondition) { return expression.containsType(If.class) - && expression.containsType(BooleanLiteral.class, NullLiteral.class); + && expression.containsType(BooleanLiteral.class, NullLiteral.class) + && !requiresShortCircuitEvaluation(expression); } @Override public Expression visitIf(If ifExpr, Boolean isInsideCondition) { + if (ifExpr instanceof RequiresShortCircuitEvaluation) { + return ifExpr; + } If newIf = (If) super.visitIf(ifExpr, isInsideCondition); if (isInsideCondition) { Expression newCondition = newIf.getCondition(); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/ConditionRewrite.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/ConditionRewrite.java index c28a395a97bd05..f97ae1f7f78560 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/ConditionRewrite.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/ConditionRewrite.java @@ -173,7 +173,7 @@ public Expression visitIf(If ifExpr, Boolean isInsideCondition) { if (newCondition != ifExpr.getCondition() || newTrueValue != ifExpr.getTrueValue() || newFalseValue != ifExpr.getFalseValue()) { - return new If(newCondition, newTrueValue, newFalseValue); + return ifExpr.withChildren(ImmutableList.of(newCondition, newTrueValue, newFalseValue)); } else { return ifExpr; } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/NestedCaseWhenCondToLiteral.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/NestedCaseWhenCondToLiteral.java index 29ddf97e786db0..ca6994af752038 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/NestedCaseWhenCondToLiteral.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/expression/rules/NestedCaseWhenCondToLiteral.java @@ -193,7 +193,7 @@ public Expression visitIf(If ifExpr, Void context) { if (newCondition != oldCondition || newTrueValue != ifExpr.getTrueValue() || newFalseValue != ifExpr.getFalseValue()) { - return new If(newCondition, newTrueValue, newFalseValue); + return ifExpr.withChildren(ImmutableList.of(newCondition, newTrueValue, newFalseValue)); } else { return ifExpr; } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/implementation/LogicalConnectorTableSinkToPhysicalConnectorTableSink.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/implementation/LogicalConnectorTableSinkToPhysicalConnectorTableSink.java index ff32bddfa139a6..3df14e08eb8094 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/implementation/LogicalConnectorTableSinkToPhysicalConnectorTableSink.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/implementation/LogicalConnectorTableSinkToPhysicalConnectorTableSink.java @@ -46,6 +46,8 @@ public Rule build() { null, null, sink.isRewrite(), + sink.getDmlCommandType(), + sink.hasRowOperationColumn(), sink.child()); }).toRule(RuleType.LOGICAL_CONNECTOR_TABLE_SINK_TO_PHYSICAL_CONNECTOR_TABLE_SINK_RULE); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/RequiresShortCircuitEvaluation.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/RequiresShortCircuitEvaluation.java new file mode 100644 index 00000000000000..880a91dbf57b31 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/RequiresShortCircuitEvaluation.java @@ -0,0 +1,22 @@ +// 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. + +package org.apache.doris.nereids.trees.expressions.functions; + +/** Marks control-flow expressions whose unselected branches must not be evaluated. */ +public interface RequiresShortCircuitEvaluation { +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/Array.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/Array.java index 8da0b6f3b45778..2941bc10728a45 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/Array.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/Array.java @@ -69,9 +69,11 @@ public void checkLegalityBeforeTypeCoercion() { if (children.isEmpty()) { return; } - DataType firstChildType = getArgument(0).getDataType(); - if (firstChildType.isJsonType() || firstChildType.isVariantType()) { - throw new AnalysisException("array does not support jsonb/variant type"); + for (Expression argument : getArguments()) { + DataType childType = argument.getDataType(); + if (childType.isJsonType()) { + throw new AnalysisException("array does not support jsonb type"); + } } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/CreateMap.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/CreateMap.java index 6bca4dbda42557..6fdcccf402b91e 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/CreateMap.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/CreateMap.java @@ -85,11 +85,16 @@ public void checkLegalityBeforeTypeCoercion() { if (arity() % 2 != 0) { throw new AnalysisException("map can't be odd parameters, need even parameters " + this.toSql()); } - children.forEach(child -> { - if (child.getDataType().isJsonType() || child.getDataType().isVariantType()) { - throw new AnalysisException("map does not support jsonb/variant type"); + for (int i = 0; i < arity(); i++) { + DataType childType = getArgument(i).getDataType(); + boolean isKey = i % 2 == 0; + if (childType.isJsonType()) { + throw new AnalysisException("map does not support jsonb type"); } - }); + if (isKey && childType.isVariantType()) { + throw new AnalysisException("map does not support variant keys"); + } + } } /** diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/CreateNamedStruct.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/CreateNamedStruct.java index 49d089bd4a696b..97164970c29830 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/CreateNamedStruct.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/CreateNamedStruct.java @@ -79,9 +79,9 @@ public void checkLegalityBeforeTypeCoercion() { names.add(name); } } - // i+1 is value, check if it is not jsonb/variant type - if (child(i + 1).getDataType().isJsonType() || child(i + 1).getDataType().isVariantType()) { - throw new AnalysisException("named_struct does not support jsonb/variant type"); + DataType valueType = getArgument(i + 1).getDataType(); + if (valueType.isJsonType()) { + throw new AnalysisException("named_struct does not support jsonb type"); } } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/CreateStruct.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/CreateStruct.java index 9e89da0fd87319..422f0e1daba583 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/CreateStruct.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/CreateStruct.java @@ -59,10 +59,9 @@ public void checkLegalityBeforeTypeCoercion() { if (arity() == 0) { throw new AnalysisException("struct requires at least one argument, like: struct(1)"); } - // for all field we do not support struct field with jsonb/variant type - children.forEach(child -> { - if (child.getDataType().isJsonType() || child.getDataType().isVariantType()) { - throw new AnalysisException("struct does not support jsonb/variant type"); + children.forEach(argument -> { + if (argument.getDataType().isJsonType()) { + throw new AnalysisException("struct does not support jsonb type"); } }); } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/ShortCircuitIf.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/ShortCircuitIf.java new file mode 100644 index 00000000000000..f019325ea3fb85 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/ShortCircuitIf.java @@ -0,0 +1,38 @@ +// 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. + +package org.apache.doris.nereids.trees.expressions.functions.scalar; + +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.functions.RequiresShortCircuitEvaluation; + +import com.google.common.base.Preconditions; + +import java.util.List; + +/** IF expression with statement-independent short-circuit semantics. */ +public class ShortCircuitIf extends If implements RequiresShortCircuitEvaluation { + public ShortCircuitIf(Expression condition, Expression trueValue, Expression falseValue) { + super(condition, trueValue, falseValue); + } + + @Override + public ShortCircuitIf withChildren(List children) { + Preconditions.checkArgument(children.size() == 3); + return new ShortCircuitIf(children.get(0), children.get(1), children.get(2)); + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ChangelogRowLevelDmlTransform.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ChangelogRowLevelDmlTransform.java new file mode 100644 index 00000000000000..1c6de60b12febb --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ChangelogRowLevelDmlTransform.java @@ -0,0 +1,229 @@ +// 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. + +package org.apache.doris.nereids.trees.plans.commands; + +import org.apache.doris.analysis.UserIdentity; +import org.apache.doris.catalog.Column; +import org.apache.doris.catalog.DatabaseIf; +import org.apache.doris.catalog.TableIf; +import org.apache.doris.connector.spi.DorisConnectorException; +import org.apache.doris.connector.spi.handle.WriteOperation; +import org.apache.doris.connector.spi.pushdown.ConnectorPredicate; +import org.apache.doris.connector.spi.write.ConnectorRowChangeStyle; +import org.apache.doris.connector.spi.write.ConnectorRowLevelDmlRequest; +import org.apache.doris.datasource.CatalogIf; +import org.apache.doris.datasource.plugin.PluginDrivenExternalTable; +import org.apache.doris.mysql.privilege.AccessControllerManager; +import org.apache.doris.nereids.NereidsPlanner; +import org.apache.doris.nereids.analyzer.UnboundConnectorTableSink; +import org.apache.doris.nereids.analyzer.UnboundRelation; +import org.apache.doris.nereids.analyzer.UnboundSlot; +import org.apache.doris.nereids.exceptions.AnalysisException; +import org.apache.doris.nereids.parser.LogicalPlanBuilderAssistant; +import org.apache.doris.nereids.trees.expressions.EqualTo; +import org.apache.doris.nereids.trees.expressions.StatementScopeIdGenerator; +import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.commands.info.ConnectorChangelogRowChangeSpec; +import org.apache.doris.nereids.trees.plans.commands.insert.BaseExternalTableInsertExecutor; +import org.apache.doris.nereids.trees.plans.commands.insert.PluginDrivenInsertExecutor; +import org.apache.doris.nereids.trees.plans.commands.merge.MergeMatchedClause; +import org.apache.doris.nereids.trees.plans.commands.merge.MergeUtils; +import org.apache.doris.nereids.trees.plans.logical.LogicalPlan; +import org.apache.doris.nereids.trees.plans.logical.LogicalSubQueryAlias; +import org.apache.doris.nereids.trees.plans.physical.PhysicalConnectorTableSink; +import org.apache.doris.nereids.trees.plans.physical.PhysicalSink; +import org.apache.doris.nereids.util.RelationUtil; +import org.apache.doris.planner.DataSink; +import org.apache.doris.planner.PlanFragment; +import org.apache.doris.qe.ConnectContext; + +import com.google.common.collect.ImmutableList; + +import java.util.List; +import java.util.Optional; +import java.util.Set; +import java.util.TreeSet; + +/** Plans row-level changes as an operation column followed by a complete table row. */ +public class ChangelogRowLevelDmlTransform implements RowLevelDmlTransform { + + @Override + public boolean handles(TableIf table) { + if (!(table instanceof PluginDrivenExternalTable)) { + return false; + } + PluginDrivenExternalTable connectorTable = (PluginDrivenExternalTable) table; + if (connectorTable.getConnectorRowChangeStyle() != ConnectorRowChangeStyle.CHANGELOG) { + return false; + } + return RowLevelDmlRegistry.supportsAnyRowLevelDml( + connectorTable.connectorSupportedWriteOperations()); + } + + @Override + public void checkMode(TableIf table, RowLevelDmlOp op) { + WriteOperation operation = op.toWriteOperation(); + if (!((PluginDrivenExternalTable) table).connectorSupportedWriteOperations().contains(operation)) { + throw new AnalysisException("Connector does not support " + operation + " operations"); + } + // Statement-specific validation runs in synthesize, where assignments and MERGE clauses are available. + } + + @Override + public LogicalPlan synthesize(ConnectContext ctx, RowLevelDmlArgs args, RowLevelDmlOp op) { + PluginDrivenExternalTable table = (PluginDrivenExternalTable) args.getTable(); + if (op == RowLevelDmlOp.DELETE && (args.isTempPart() || !args.getPartitions().isEmpty())) { + throw new AnalysisException( + "Connector changelog DELETE does not support partition name lists; use a WHERE predicate"); + } + validate(ctx, table, args, op); + switch (op) { + case DELETE: + return deletePlan(ctx, args); + case UPDATE: + return updatePlan(ctx, args); + default: + return mergePlan(ctx, args); + } + } + + private LogicalPlan deletePlan(ConnectContext ctx, RowLevelDmlArgs args) { + List target = args.getTableAlias() != null + ? ImmutableList.of(args.getTableAlias()) + : RelationUtil.getQualifierName(ctx, args.getNameParts()); + return new UnboundConnectorTableSink<>(args.getNameParts(), args.getLogicalQuery(), + new ConnectorChangelogRowChangeSpec.Delete(target, args.shouldDeduplicateTargetRows())); + } + + private LogicalPlan updatePlan(ConnectContext ctx, RowLevelDmlArgs args) { + for (EqualTo assignment : args.getAssignments()) { + UpdateCommand.checkAssignmentColumn(ctx, + ((UnboundSlot) assignment.left()).getNameParts(), + args.getNameParts(), args.getTableAlias()); + } + List target = args.getTableAlias() != null + ? ImmutableList.of(args.getTableAlias()) + : RelationUtil.getQualifierName(ctx, args.getNameParts()); + LogicalPlan sink = new UnboundConnectorTableSink<>(args.getNameParts(), args.getLogicalQuery(), + new ConnectorChangelogRowChangeSpec.Update(target, args.getAssignments())); + return args.getCte().isPresent() ? (LogicalPlan) args.getCte().get().withChildren(sink) : sink; + } + + private LogicalPlan mergePlan(ConnectContext ctx, RowLevelDmlArgs args) { + for (MergeMatchedClause clause : args.getMatchedClauses()) { + for (EqualTo assignment : clause.getAssignments()) { + UpdateCommand.checkAssignmentColumn(ctx, + ((UnboundSlot) assignment.left()).getNameParts(), + args.getTargetNameParts(), args.getTargetAlias().orElse(null)); + } + } + List targetName = args.getTargetAlias().isPresent() + ? ImmutableList.of(args.getTargetAlias().get()) + : RelationUtil.getQualifierName(ctx, args.getTargetNameParts()); + ConnectorChangelogRowChangeSpec.Merge spec = new ConnectorChangelogRowChangeSpec.Merge( + targetName, args.getMatchedClauses(), args.getNotMatchedClauses()); + LogicalPlan target = LogicalPlanBuilderAssistant.withCheckPolicy( + new UnboundRelation(StatementScopeIdGenerator.newRelationId(), args.getTargetNameParts())); + if (args.getTargetAlias().isPresent()) { + target = new LogicalSubQueryAlias<>(args.getTargetAlias().get(), target); + } + LogicalPlan join = MergeUtils.buildMergeJoin(target, args.getSource(), args.getOnClause(), + !args.getNotMatchedClauses().isEmpty()); + LogicalPlan sink = new UnboundConnectorTableSink<>(args.getTargetNameParts(), join, spec); + return args.getCte().isPresent() ? (LogicalPlan) args.getCte().get().withChildren(sink) : sink; + } + + private void validate(ConnectContext ctx, PluginDrivenExternalTable table, + RowLevelDmlArgs args, RowLevelDmlOp op) { + requireNoDataMask(ctx, table, op); + Set updatedColumns = new TreeSet<>(String.CASE_INSENSITIVE_ORDER); + boolean containsUpdate = op == RowLevelDmlOp.UPDATE; + boolean containsDelete = op == RowLevelDmlOp.DELETE; + if (op == RowLevelDmlOp.UPDATE) { + addUpdatedColumns(updatedColumns, args.getAssignments()); + } else if (op == RowLevelDmlOp.MERGE) { + for (MergeMatchedClause clause : args.getMatchedClauses()) { + containsDelete |= clause.isDelete(); + containsUpdate |= !clause.isDelete(); + addUpdatedColumns(updatedColumns, clause.getAssignments()); + } + } + try { + table.validateConnectorRowLevelDml(new ConnectorRowLevelDmlRequest( + op.toWriteOperation(), updatedColumns, containsUpdate, containsDelete)); + } catch (DorisConnectorException e) { + throw new AnalysisException(e.getMessage(), e); + } + } + + static void requireNoDataMask(ConnectContext ctx, PluginDrivenExternalTable table, RowLevelDmlOp op) { + UserIdentity user = ctx.getCurrentUserIdentity(); + if (user.isRootUser() || user.isAdminUser()) { + return; + } + DatabaseIf database = table.getDatabase(); + CatalogIf catalog = database.getCatalog(); + Set columns = new TreeSet<>(String.CASE_INSENSITIVE_ORDER); + for (Column column : table.getFullSchema()) { + columns.add(column.getName()); + } + AccessControllerManager accessManager = ctx.getEnv().getAccessManager(); + if (!accessManager.evalDataMaskPolicies( + user, catalog.getName(), database.getFullName(), table.getName(), columns).isEmpty()) { + throw new AnalysisException("Connector " + op + + " is not supported when data masking policies apply to the target table"); + } + } + + private void addUpdatedColumns(Set columns, List assignments) { + for (EqualTo assignment : assignments) { + List parts = ((UnboundSlot) assignment.left()).getNameParts(); + columns.add(parts.get(parts.size() - 1)); + } + } + + @Override + public BaseExternalTableInsertExecutor newExecutor(ConnectContext ctx, TableIf table, String label, + NereidsPlanner planner, boolean emptyInsert, RowLevelDmlOp op) { + return new PluginDrivenInsertExecutor(ctx, (PluginDrivenExternalTable) table, label, + planner, Optional.empty(), emptyInsert, -1L); + } + + @Override + public PhysicalSink requirePhysicalSink(NereidsPlanner planner, RowLevelDmlOp op) { + return planner.getPhysicalPlan().>collect(PhysicalSink.class::isInstance) + .stream().filter(PhysicalConnectorTableSink.class::isInstance).findAny() + .orElseThrow(() -> new AnalysisException(op + " plan must use connector table sink")); + } + + @Override + public String labelPrefix(TableIf table, RowLevelDmlOp op) { + return ((PluginDrivenExternalTable) table).getConnectorRowLevelDmlLabelPrefix(op.toWriteOperation()); + } + + @Override + public void finalizeSink(BaseExternalTableInsertExecutor executor, RowLevelDmlOp op, + PlanFragment fragment, DataSink sink, PhysicalSink physicalSink) { + ((PluginDrivenInsertExecutor) executor).finalizeRowLevelDmlSink(fragment, sink, physicalSink); + } + + @Override + public Optional extractWriteConstraint(Plan analyzedPlan, TableIf table) { + return Optional.empty(); + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ConnectorWriteSchemaUtils.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ConnectorWriteSchemaUtils.java index f3efe1d3115b78..0fe40e06509458 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ConnectorWriteSchemaUtils.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ConnectorWriteSchemaUtils.java @@ -37,7 +37,7 @@ import java.util.Optional; /** Neutral engine helpers for request-scoped connector writer schemas and DEFAULT expressions. */ -final class ConnectorWriteSchemaUtils { +public final class ConnectorWriteSchemaUtils { private ConnectorWriteSchemaUtils() { } @@ -104,7 +104,8 @@ static Expression resolveDefaultReferences(Expression expression, List c }); } - static Expression resolveDefault(Column column) { + /** Parse the catalog default expression for a connector write column. */ + public static Expression resolveDefault(Column column) { String defaultSql = column.getDefaultValueSql(); if (defaultSql == null) { throw new AnalysisException( @@ -114,7 +115,8 @@ static Expression resolveDefault(Column column) { return expression instanceof UnboundAlias ? expression.child(0) : expression; } - static Expression resolveExplicitDefault(Expression expression, Column column) { + /** Replace an explicit DEFAULT placeholder with the column's catalog default expression. */ + public static Expression resolveExplicitDefault(Expression expression, Column column) { return expression instanceof DefaultValueSlot ? resolveDefault(column) : expression; } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/DeleteFromCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/DeleteFromCommand.java index bda236a15bcd0d..4a2468f613f265 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/DeleteFromCommand.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/DeleteFromCommand.java @@ -142,8 +142,7 @@ public void run(ConnectContext ctx, StmtExecutor executor) throws Exception { // Route row-level DML on external tables (e.g. iceberg) through the generic shell. Optional transform = RowLevelDmlRegistry.find(table); if (transform.isPresent()) { - RowLevelDmlArgs args = RowLevelDmlArgs.forDelete( - table, nameParts, tableAlias, isTempPart, partitions, logicalQuery); + RowLevelDmlArgs args = rowLevelDmlArgs(table); new RowLevelDmlCommand(transform.get(), args, RowLevelDmlOp.DELETE).run(ctx, executor); return; } @@ -499,13 +498,17 @@ public Plan getExplainPlan(ConnectContext ctx) { TableIf table = RelationUtil.getTable(qualifiedTableName, ctx.getEnv(), Optional.empty()); Optional transform = RowLevelDmlRegistry.find(table); if (transform.isPresent()) { - RowLevelDmlArgs args = RowLevelDmlArgs.forDelete( - table, nameParts, tableAlias, isTempPart, partitions, logicalQuery); + RowLevelDmlArgs args = rowLevelDmlArgs(table); return new RowLevelDmlCommand(transform.get(), args, RowLevelDmlOp.DELETE).getExplainPlan(ctx); } return completeQueryPlan(ctx, logicalQuery); } + protected RowLevelDmlArgs rowLevelDmlArgs(TableIf table) { + return RowLevelDmlArgs.forDelete( + table, nameParts, tableAlias, isTempPart, partitions, logicalQuery); + } + private OlapTable getTargetTable(ConnectContext ctx) { List qualifiedTableName = RelationUtil.getQualifierName(ctx, nameParts); TableIf table = RelationUtil.getTable(qualifiedTableName, ctx.getEnv(), Optional.empty()); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/DeleteFromUsingCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/DeleteFromUsingCommand.java index 6d6b53d2254a56..4d16e5d3c33525 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/DeleteFromUsingCommand.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/DeleteFromUsingCommand.java @@ -20,10 +20,12 @@ import org.apache.doris.analysis.StmtType; import org.apache.doris.catalog.KeysType; import org.apache.doris.catalog.OlapTable; +import org.apache.doris.catalog.TableIf; import org.apache.doris.nereids.exceptions.AnalysisException; import org.apache.doris.nereids.trees.plans.commands.insert.InsertIntoTableCommand; import org.apache.doris.nereids.trees.plans.logical.LogicalPlan; import org.apache.doris.nereids.trees.plans.visitor.PlanVisitor; +import org.apache.doris.nereids.util.RelationUtil; import org.apache.doris.qe.ConnectContext; import org.apache.doris.qe.StmtExecutor; @@ -55,6 +57,14 @@ public void run(ConnectContext ctx, StmtExecutor executor) throws Exception { + " Please check the following session variables: " + ctx.getSessionVariable().printDebugModeVariables()); } + TableIf table = RelationUtil.getTable(RelationUtil.getQualifierName(ctx, nameParts), + ctx.getEnv(), Optional.empty()); + Optional transform = RowLevelDmlRegistry.find(table); + if (transform.isPresent()) { + RowLevelDmlArgs args = rowLevelDmlArgs(table); + new RowLevelDmlCommand(transform.get(), args, RowLevelDmlOp.DELETE).run(ctx, executor); + return; + } // NOTE: delete from using command is executed as insert command, so txn insert can support it new InsertIntoTableCommand(completeQueryPlan(ctx, logicalQuery), Optional.empty(), Optional.empty(), Optional.empty(), true, Optional.empty()).run(ctx, executor); @@ -68,6 +78,12 @@ protected LogicalPlan handleCte(LogicalPlan logicalPlan) { return logicalPlan; } + @Override + protected RowLevelDmlArgs rowLevelDmlArgs(TableIf table) { + return RowLevelDmlArgs.forDelete( + table, nameParts, tableAlias, isTempPart, partitions, handleCte(logicalQuery), true); + } + /** * for test */ diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExternalRowLevelDeletePlanBuilder.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExternalRowLevelDeletePlanBuilder.java index 1decb589794a9d..38af844ad87ad9 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExternalRowLevelDeletePlanBuilder.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExternalRowLevelDeletePlanBuilder.java @@ -40,7 +40,7 @@ /** * DELETE plan synthesizer for Iceberg tables, invoked by - * IcebergRowLevelDmlTransform.synthesize via {@link #completeQueryPlan}. + * PositionDeleteRowLevelDmlTransform.synthesize via {@link #completeQueryPlan}. * * It rewrites a DELETE into an insert-shaped plan that generates * position DeleteFile entries instead of data files. diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExternalRowLevelMergePlanBuilder.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExternalRowLevelMergePlanBuilder.java index fa92a2922a3692..82be8acd764482 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExternalRowLevelMergePlanBuilder.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExternalRowLevelMergePlanBuilder.java @@ -37,7 +37,6 @@ import org.apache.doris.nereids.trees.expressions.Not; import org.apache.doris.nereids.trees.expressions.Slot; import org.apache.doris.nereids.trees.expressions.StatementScopeIdGenerator; -import org.apache.doris.nereids.trees.expressions.functions.scalar.If; import org.apache.doris.nereids.trees.expressions.literal.IntegerLiteral; import org.apache.doris.nereids.trees.expressions.literal.NullLiteral; import org.apache.doris.nereids.trees.expressions.literal.TinyIntLiteral; @@ -71,7 +70,7 @@ import java.util.Optional; /** - * Iceberg MERGE INTO plan synthesizer, invoked via IcebergRowLevelDmlTransform.synthesize + * Position-delete MERGE INTO plan synthesizer, invoked via PositionDeleteRowLevelDmlTransform.synthesize * (legacy execution half removed as dead code). */ public class ExternalRowLevelMergePlanBuilder { @@ -135,7 +134,8 @@ private NamedExpression generateBranchLabel(Expression rowIdExpr) { } Expression currentResult = new IntegerLiteral(i); if (clause.getCasePredicate().isPresent()) { - matchedLabel = new If(clause.getCasePredicate().get(), currentResult, matchedLabel); + matchedLabel = MergeUtils.selectBranch( + clause.getCasePredicate().get(), currentResult, matchedLabel); } else { matchedLabel = currentResult; } @@ -149,13 +149,15 @@ private NamedExpression generateBranchLabel(Expression rowIdExpr) { } Expression currentResult = new IntegerLiteral(i + matchedClauses.size()); if (clause.getCasePredicate().isPresent()) { - notMatchedLabel = new If(clause.getCasePredicate().get(), currentResult, notMatchedLabel); + notMatchedLabel = MergeUtils.selectBranch( + clause.getCasePredicate().get(), currentResult, notMatchedLabel); } else { notMatchedLabel = currentResult; } } - return new UnboundAlias(new If(new Not(new IsNull(rowIdExpr)), matchedLabel, notMatchedLabel), + return new UnboundAlias(MergeUtils.selectBranch( + new Not(new IsNull(rowIdExpr)), matchedLabel, notMatchedLabel), BRANCH_LABEL); } @@ -308,7 +310,8 @@ static List generateFinalProjections(List colNames, for (int j = 0; j < finalProjections.size(); j++) { Expression branch = TypeCoercionUtils.castUnbound( finalProjections.get(j).get(i), outputType); - project = new If(new EqualTo(new UnboundSlot(BRANCH_LABEL), new IntegerLiteral(j)), + project = MergeUtils.selectBranch( + new EqualTo(new UnboundSlot(BRANCH_LABEL), new IntegerLiteral(j)), branch, project); } output.add(new UnboundAlias(project, colNames.get(i))); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExternalRowLevelUpdatePlanBuilder.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExternalRowLevelUpdatePlanBuilder.java index 240304ea458b74..0319fe2e2212a0 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExternalRowLevelUpdatePlanBuilder.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/ExternalRowLevelUpdatePlanBuilder.java @@ -49,7 +49,7 @@ /** * Merge-plan synthesizer for UPDATE on Iceberg tables, invoked via - * IcebergRowLevelDmlTransform.synthesize. The legacy Command execution half + * PositionDeleteRowLevelDmlTransform.synthesize. The legacy Command execution half * was removed as dead code. * * UPDATE operations are implemented as a single scan + merge sink: diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/IcebergRowLevelDmlTransform.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/PositionDeleteRowLevelDmlTransform.java similarity index 67% rename from fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/IcebergRowLevelDmlTransform.java rename to fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/PositionDeleteRowLevelDmlTransform.java index 5e9ba81668f136..9996e3e891539e 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/IcebergRowLevelDmlTransform.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/PositionDeleteRowLevelDmlTransform.java @@ -17,7 +17,6 @@ package org.apache.doris.nereids.trees.plans.commands; -import org.apache.doris.catalog.Column; import org.apache.doris.catalog.TableIf; import org.apache.doris.connector.spi.ConnectorMetadata; import org.apache.doris.connector.spi.ConnectorSession; @@ -25,6 +24,7 @@ import org.apache.doris.connector.spi.handle.ConnectorTableHandle; import org.apache.doris.connector.spi.handle.WriteOperation; import org.apache.doris.connector.spi.pushdown.ConnectorPredicate; +import org.apache.doris.connector.spi.write.ConnectorRowChangeStyle; import org.apache.doris.datasource.ExternalTable; import org.apache.doris.datasource.connector.converter.WriteConstraintExtractor; import org.apache.doris.datasource.plugin.PluginDrivenExternalCatalog; @@ -44,51 +44,34 @@ import org.apache.doris.planner.PlanFragment; import org.apache.doris.qe.ConnectContext; -import com.google.common.collect.ImmutableSet; - import java.util.Optional; import java.util.Set; +import java.util.TreeSet; import java.util.function.Predicate; /** - * Iceberg {@link RowLevelDmlTransform}: routes {@code DELETE}/{@code UPDATE}/{@code MERGE INTO} on iceberg - * tables through the generic {@link RowLevelDmlCommand} shell. + * Position-delete {@link RowLevelDmlTransform}: routes {@code DELETE}/{@code UPDATE}/{@code MERGE INTO} + * through the generic {@link RowLevelDmlCommand} shell. * - *

The iceberg plan-synthesis algebra lives in same-package neutral helpers: {@link #synthesize} constructs + *

The plan-synthesis algebra lives in same-package neutral helpers: {@link #synthesize} constructs * the corresponding {@code ExternalRowLevel*PlanBuilder} and calls its (package-visible) synthesis method, so * the synthesized {@code LogicalExternalRowLevel{Delete,Merge}Sink} tree is the generic row-level DML sink. * The per-executor-only bits (conflict-filter stash, finalize) are routed here via - * {@code instanceof}-free op switches; the exclusion predicate mirrors legacy - * {@code IcebergConflictDetectionFilterUtils} (note the {@code equalsIgnoreCase} vs {@code equals} asymmetry).

+ * {@code instanceof}-free operation switches. Connector-owned metadata column names are obtained through + * the write-provider SPI rather than embedded in engine code.

*/ -public class IcebergRowLevelDmlTransform implements RowLevelDmlTransform { - - /** - * Position-delete metadata column names ({@code $file_path}/{@code $row_position}/{@code $partition_spec_id}/ - * {@code $partition_data}): the connector-declared row-id STRUCT field names, {@code $}-prefixed. Kept as - * FE-side synthetic-column name constants (the same category as {@link Column#ICEBERG_ROWID_COL}); matched - * case-sensitively ({@code equals}), unlike the rowid ({@code equalsIgnoreCase}). - */ - private static final Set ICEBERG_METADATA_COLUMN_NAMES = ImmutableSet.of( - "$file_path", "$row_position", "$partition_spec_id", "$partition_data"); - - /** - * Slots excluded from the target-only write constraint: the synthetic {@code $row_id} column and - * iceberg metadata columns. Mirrors legacy {@code IcebergConflictDetectionFilterUtils.isTargetOnlyPredicate} - * exactly — keep the {@code equalsIgnoreCase} (rowid) vs {@code equals} (metadata) asymmetry. - */ - private static final Predicate ICEBERG_EXCLUSION = - slot -> Column.ICEBERG_ROWID_COL.equalsIgnoreCase(slot.getName()) - || ICEBERG_METADATA_COLUMN_NAMES.contains(slot.getName()); +public class PositionDeleteRowLevelDmlTransform implements RowLevelDmlTransform { @Override public boolean handles(TableIf table) { return table instanceof PluginDrivenExternalTable + && ((PluginDrivenExternalTable) table).getConnectorRowChangeStyle() + == ConnectorRowChangeStyle.POSITION_DELETE && pluginConnectorSupportsRowLevelDml((PluginDrivenExternalTable) table); } /** - * A plugin-driven (SPI connector) table is routed through the iceberg row-level DML synthesis only if + * A plugin-driven table is routed through position-delete row-level DML synthesis only if * its connector declares row-level DML support ({@code supportsDelete()} or {@code supportsMerge()}). * Mirrors the connector-capability probe in * {@code InsertOverwriteTableCommand.pluginConnectorSupportsInsertOverwrite}. @@ -97,28 +80,33 @@ public boolean handles(TableIf table) { * admits "supports any row-level DML"; per-op validity (e.g. UPDATE against a delete-only connector) is * enforced later in {@link #checkMode}.

* - *

Today only the iceberg connector declares these capabilities (every other SPI connector inherits - * the {@code ConnectorWriteOps} default {@code false}).

+ *

The representation check in {@link #handles} must precede this capability check: a connector + * using changelog rows may support the same operations but cannot use the position-delete plan.

*/ private static boolean pluginConnectorSupportsRowLevelDml(PluginDrivenExternalTable table) { - // Per-handle write-op probe: a heterogeneous gateway admits row-level DML for its iceberg tables only. + // Per-handle write-op probe lets a heterogeneous gateway select only qualifying tables. Set ops = table.connectorSupportedWriteOperations(); - return ops.contains(WriteOperation.DELETE) || ops.contains(WriteOperation.MERGE); + return RowLevelDmlRegistry.supportsAnyRowLevelDml(ops); } @Override public void checkMode(TableIf table, RowLevelDmlOp op) { - checkPluginMode((PluginDrivenExternalTable) table, op); + PluginDrivenExternalTable connectorTable = (PluginDrivenExternalTable) table; + WriteOperation operation = op.toWriteOperation(); + if (!connectorTable.connectorSupportedWriteOperations().contains(operation)) { + throw new AnalysisException("Connector does not support " + operation + " operations"); + } + checkPluginMode(connectorTable, operation); } /** * {@link #checkMode} body: route the copy-on-write rejection through the connector's neutral - * {@code validateRowLevelDmlMode} SPI, so the iceberg property knowledge and the message stay in the + * {@code validateRowLevelDmlMode} SPI, so format-specific properties and messages stay in the * connector. A connector {@link DorisConnectorException} is surfaced as the analysis-time * {@link AnalysisException} the legacy native path threw, preserving the user-facing message and the * exception type. */ - private static void checkPluginMode(PluginDrivenExternalTable table, RowLevelDmlOp op) { + private static void checkPluginMode(PluginDrivenExternalTable table, WriteOperation operation) { PluginDrivenExternalCatalog catalog = (PluginDrivenExternalCatalog) table.getCatalog(); ConnectorSession session = catalog.buildConnectorSession(); ConnectorMetadata metadata = PluginDrivenMetadata.get(session, catalog.getConnector()); @@ -128,42 +116,31 @@ private static void checkPluginMode(PluginDrivenExternalTable table, RowLevelDml + table.getRemoteDbName() + "." + table.getRemoteName() + " in catalog " + catalog.getName())); try { - metadata.validateRowLevelDmlMode(session, handle, toWriteOperation(op)); + metadata.validateRowLevelDmlMode(session, handle, operation); } catch (DorisConnectorException e) { throw new AnalysisException(e.getMessage(), e); } } - private static WriteOperation toWriteOperation(RowLevelDmlOp op) { - switch (op) { - case DELETE: - return WriteOperation.DELETE; - case UPDATE: - return WriteOperation.UPDATE; - default: - return WriteOperation.MERGE; - } - } - @Override public LogicalPlan synthesize(ConnectContext ctx, RowLevelDmlArgs args, RowLevelDmlOp op) { - ExternalTable icebergTable = (ExternalTable) args.getTable(); + ExternalTable externalTable = (ExternalTable) args.getTable(); switch (op) { case DELETE: return new ExternalRowLevelDeletePlanBuilder( args.getNameParts(), args.getTableAlias(), args.isTempPart(), args.getPartitions(), args.getLogicalQuery()) - .completeQueryPlan(ctx, args.getLogicalQuery(), icebergTable); + .completeQueryPlan(ctx, args.getLogicalQuery(), externalTable); case UPDATE: return new ExternalRowLevelUpdatePlanBuilder( args.getNameParts(), args.getTableAlias(), args.getAssignments(), args.getLogicalQuery()) - .buildMergePlan(ctx, args.getLogicalQuery(), args.getAssignments(), icebergTable); + .buildMergePlan(ctx, args.getLogicalQuery(), args.getAssignments(), externalTable); default: return new ExternalRowLevelMergePlanBuilder( args.getTargetNameParts(), args.getTargetAlias(), args.getCte(), args.getSource(), args.getOnClause(), args.getMatchedClauses(), args.getNotMatchedClauses()) - .buildMergePlan(ctx, icebergTable); + .buildMergePlan(ctx, externalTable); } } @@ -188,7 +165,7 @@ public PhysicalSink requirePhysicalSink(NereidsPlanner planner, RowLevelDmlOp throw new AnalysisException("DELETE command must contain target table"); } if (!(plan.get() instanceof PhysicalExternalRowLevelDeleteSink)) { - throw new AnalysisException("DELETE plan must use Iceberg delete sink"); + throw new AnalysisException("DELETE plan must use a position-delete sink"); } return plan.get(); case UPDATE: @@ -196,7 +173,7 @@ public PhysicalSink requirePhysicalSink(NereidsPlanner planner, RowLevelDmlOp throw new AnalysisException("UPDATE command must contain target table"); } if (!(plan.get() instanceof PhysicalExternalRowLevelMergeSink)) { - throw new AnalysisException("UPDATE merge plan must use Iceberg merge sink"); + throw new AnalysisException("UPDATE plan must use a position-delete merge sink"); } return plan.get(); default: @@ -204,33 +181,21 @@ public PhysicalSink requirePhysicalSink(NereidsPlanner planner, RowLevelDmlOp throw new AnalysisException("MERGE INTO command must contain target table"); } if (!(plan.get() instanceof PhysicalExternalRowLevelMergeSink)) { - throw new AnalysisException("MERGE INTO plan must use Iceberg merge sink"); + throw new AnalysisException("MERGE INTO plan must use a position-delete merge sink"); } return plan.get(); } } @Override - public String labelPrefix(RowLevelDmlOp op) { - switch (op) { - case DELETE: - return "iceberg_delete"; - case UPDATE: - return "iceberg_update_merge"; - default: - return "iceberg_merge_into"; - } + public String labelPrefix(TableIf table, RowLevelDmlOp op) { + return ((PluginDrivenExternalTable) table) + .getConnectorRowLevelDmlLabelPrefix(op.toWriteOperation()); } @Override - public void setupConflictDetection(BaseExternalTableInsertExecutor executor, Plan analyzedPlan, TableIf table, - RowLevelDmlOp op) { - // No-op: the conflict filter is supplied through the neutral SPI path - // (RowLevelDmlCommand.applyWriteConstraintIfPresent -> extractWriteConstraint -> - // ConnectorTransaction.applyWriteConstraint), converted to a native iceberg Expression lazily at - // commit. Running ONLY the SPI path avoids double-filtering; the SPI converter is byte-verified - // equivalent to the retired native filter builder, the residual divergence only widening the - // filter -> at worst a harmless extra OCC retry (see [DEC-S5]). + public boolean requiresExternalTableBatchModeDisabled() { + return true; } @Override @@ -244,6 +209,10 @@ public void finalizeSink(BaseExternalTableInsertExecutor executor, RowLevelDmlOp @Override public Optional extractWriteConstraint(Plan analyzedPlan, TableIf table) { - return WriteConstraintExtractor.extract(analyzedPlan, table.getId(), ICEBERG_EXCLUSION); + Set excludedColumns = new TreeSet<>(String.CASE_INSENSITIVE_ORDER); + excludedColumns.addAll(((PluginDrivenExternalTable) table) + .getConnectorRowLevelWriteConstraintExcludedColumns()); + Predicate exclusion = slot -> excludedColumns.contains(slot.getName()); + return WriteConstraintExtractor.extract(analyzedPlan, table.getId(), exclusion); } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/RowLevelDmlArgs.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/RowLevelDmlArgs.java index c77b2e45ed4dff..0ab125d7893188 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/RowLevelDmlArgs.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/RowLevelDmlArgs.java @@ -49,6 +49,7 @@ public final class RowLevelDmlArgs { // DELETE only private final boolean isTempPart; private final List partitions; + private final boolean deduplicateTargetRows; // UPDATE only private final List assignments; @@ -63,6 +64,7 @@ public final class RowLevelDmlArgs { private RowLevelDmlArgs(TableIf table, List nameParts, String tableAlias, LogicalPlan logicalQuery, boolean isTempPart, List partitions, List assignments, + boolean deduplicateTargetRows, List targetNameParts, Optional targetAlias, Optional cte, LogicalPlan source, Expression onClause, List matchedClauses, List notMatchedClauses) { @@ -72,6 +74,7 @@ private RowLevelDmlArgs(TableIf table, List nameParts, String tableAlias this.logicalQuery = logicalQuery; this.isTempPart = isTempPart; this.partitions = partitions; + this.deduplicateTargetRows = deduplicateTargetRows; this.assignments = assignments; this.targetNameParts = targetNameParts; this.targetAlias = targetAlias; @@ -86,14 +89,26 @@ private RowLevelDmlArgs(TableIf table, List nameParts, String tableAlias public static RowLevelDmlArgs forDelete(TableIf table, List nameParts, String tableAlias, boolean isTempPart, List partitions, LogicalPlan logicalQuery) { return new RowLevelDmlArgs(table, nameParts, tableAlias, logicalQuery, isTempPart, partitions, - null, null, null, null, null, null, null, null); + null, false, null, null, null, null, null, null, null); + } + + public static RowLevelDmlArgs forDelete(TableIf table, List nameParts, String tableAlias, + boolean isTempPart, List partitions, LogicalPlan logicalQuery, + boolean deduplicateTargetRows) { + return new RowLevelDmlArgs(table, nameParts, tableAlias, logicalQuery, isTempPart, partitions, + null, deduplicateTargetRows, null, null, null, null, null, null, null); } /** Arguments for an UPDATE (mirrors the legacy {@code ExternalRowLevelUpdatePlanBuilder} constructor inputs). */ public static RowLevelDmlArgs forUpdate(TableIf table, List nameParts, String tableAlias, List assignments, LogicalPlan logicalQuery) { + return forUpdate(table, nameParts, tableAlias, assignments, logicalQuery, Optional.empty()); + } + + public static RowLevelDmlArgs forUpdate(TableIf table, List nameParts, String tableAlias, + List assignments, LogicalPlan logicalQuery, Optional cte) { return new RowLevelDmlArgs(table, nameParts, tableAlias, logicalQuery, false, null, - assignments, null, null, null, null, null, null, null); + assignments, false, null, null, cte, null, null, null, null); } /** Arguments for a MERGE INTO (mirrors the legacy {@code ExternalRowLevelMergePlanBuilder} constructor inputs). */ @@ -101,7 +116,7 @@ public static RowLevelDmlArgs forMerge(TableIf table, List targetNamePar Optional cte, LogicalPlan source, Expression onClause, List matchedClauses, List notMatchedClauses) { return new RowLevelDmlArgs(table, null, null, null, false, null, null, - targetNameParts, targetAlias, cte, source, onClause, matchedClauses, notMatchedClauses); + false, targetNameParts, targetAlias, cte, source, onClause, matchedClauses, notMatchedClauses); } public TableIf getTable() { @@ -128,6 +143,10 @@ public List getPartitions() { return partitions; } + public boolean shouldDeduplicateTargetRows() { + return deduplicateTargetRows; + } + public List getAssignments() { return assignments; } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/RowLevelDmlCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/RowLevelDmlCommand.java index c992f3a5a520af..e944f75ffbe69d 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/RowLevelDmlCommand.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/RowLevelDmlCommand.java @@ -42,8 +42,7 @@ * *

Owns the single live planner-drive loop that was triplicated across {@code ExternalRowLevelDeletePlanBuilder}, * {@code ExternalRowLevelUpdatePlanBuilder} and {@code ExternalRowLevelMergePlanBuilder}: the per-operation - * points (mode check, plan synthesis, required sink, executor factory, label prefix, conflict-detection - * wiring, finalize) are routed + * points (mode check, plan synthesis, required sink, executor factory, label prefix and finalize) are routed * through a {@link RowLevelDmlTransform} resolved from {@link RowLevelDmlRegistry}. The dispatching commands * ({@code UpdateCommand}/{@code DeleteFromCommand}/{@code MergeIntoCommand}) delegate here once a transform is * found, so the reverse {@code instanceof} dispatch is consolidated into the registry.

@@ -78,7 +77,7 @@ public void run(ConnectContext ctx, StmtExecutor stmtExecutor) throws Exception ctx.setSyntheticWriteColTargetTableId(table.getId()); try { LogicalPlan plan = transform.synthesize(ctx, args, op); - executeWithExternalTableBatchModeDisabled(ctx, () -> { + Callable execute = () -> { LogicalPlanAdapter logicalPlanAdapter = new LogicalPlanAdapter(plan, ctx.getStatementContext()); NereidsPlanner planner = new NereidsPlanner(ctx.getStatementContext()); planner.plan(logicalPlanAdapter, ctx.getSessionVariable().toThrift()); @@ -89,12 +88,11 @@ public void run(ConnectContext ctx, StmtExecutor stmtExecutor) throws Exception PlanFragment fragment = planner.getFragments().get(0); DataSink dataSink = fragment.getSink(); boolean emptyInsert = childIsEmptyRelation(physicalSink); - String label = String.format(transform.labelPrefix(op) + "_%x_%x", + String label = String.format(transform.labelPrefix(args.getTable(), op) + "_%x_%x", ctx.queryId().hi, ctx.queryId().lo); BaseExternalTableInsertExecutor insertExecutor = transform.newExecutor(ctx, table, label, planner, emptyInsert, op); - transform.setupConflictDetection(insertExecutor, planner.getAnalyzedPlan(), table, op); if (insertExecutor.isEmptyInsert()) { return null; @@ -104,7 +102,12 @@ public void run(ConnectContext ctx, StmtExecutor stmtExecutor) throws Exception planner.getAnalyzedPlan(), table, fragment, dataSink, physicalSink); insertExecutor.executeSingleInsert(stmtExecutor); return null; - }); + }; + if (transform.requiresExternalTableBatchModeDisabled()) { + executeWithExternalTableBatchModeDisabled(ctx, execute); + } else { + execute.call(); + } } finally { ctx.setSyntheticWriteColTargetTableId(previousTargetTableId); } @@ -149,12 +152,7 @@ static void beginTransactionAndFinalizeSink(RowLevelDmlTransform transform, RowL } } - /** - * Write-constraint path: only fires when the executor exposes an SPI - * {@link ConnectorTransaction}. Today iceberg DELETE/MERGE run on the legacy {@code IcebergTransaction} - * (the base {@code getConnectorTransactionOrNull()} returns {@code null}), so this is a no-op; the legacy - * 3-hop conflict-detection path ({@link RowLevelDmlTransform#setupConflictDetection}) remains the live one. - */ + /** Applies the connector-neutral optimistic write constraint when the transaction supports it. */ @VisibleForTesting static void applyWriteConstraintIfPresent(RowLevelDmlTransform transform, BaseExternalTableInsertExecutor executor, Plan analyzedPlan, TableIf table) { @@ -166,9 +164,8 @@ static void applyWriteConstraintIfPresent(RowLevelDmlTransform transform, } /** - * Run {@code action} with external-table batch mode disabled so the iceberg scan node yields all splits - * (needed by {@code IcebergRewritableDeletePlanner.collect}). Byte-identical to the per-command copies - * retained on the legacy {@code Iceberg*Command} classes until P6.7. + * Runs {@code action} with external-table batch mode disabled for row-change representations that require + * every source split to be available while the write is planned. */ static T executeWithExternalTableBatchModeDisabled(ConnectContext ctx, Callable action) throws Exception { boolean previousEnableExternalTableBatchMode = ctx.getSessionVariable().enableExternalTableBatchMode; diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/RowLevelDmlOp.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/RowLevelDmlOp.java index 1ed1e145780b03..913f14b959c2a4 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/RowLevelDmlOp.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/RowLevelDmlOp.java @@ -17,6 +17,8 @@ package org.apache.doris.nereids.trees.plans.commands; +import org.apache.doris.connector.spi.handle.WriteOperation; + /** * The kind of row-level DML driven by the generic {@link RowLevelDmlCommand} shell. * @@ -26,5 +28,19 @@ public enum RowLevelDmlOp { DELETE, UPDATE, - MERGE + MERGE; + + /** Returns the connector SPI operation corresponding to this command. */ + public WriteOperation toWriteOperation() { + switch (this) { + case DELETE: + return WriteOperation.DELETE; + case UPDATE: + return WriteOperation.UPDATE; + case MERGE: + return WriteOperation.MERGE; + default: + throw new IllegalStateException("Unsupported row-level DML operation: " + this); + } + } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/RowLevelDmlRegistry.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/RowLevelDmlRegistry.java index 9a2a0919ed47a8..59b38f822feb1e 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/RowLevelDmlRegistry.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/RowLevelDmlRegistry.java @@ -18,24 +18,28 @@ package org.apache.doris.nereids.trees.plans.commands; import org.apache.doris.catalog.TableIf; +import org.apache.doris.connector.spi.handle.WriteOperation; +import org.apache.doris.datasource.plugin.PluginDrivenExternalTable; +import org.apache.doris.nereids.exceptions.AnalysisException; import com.google.common.collect.ImmutableList; import java.util.List; import java.util.Optional; +import java.util.Set; /** * Registry of {@link RowLevelDmlTransform}s. The dispatching DML commands consult this instead of testing the * target table type, so the reverse {@code instanceof} dispatch is consolidated here. * *

Explicit static registration (no {@code ServiceLoader}) — avoids the thread-context-classloader pitfalls - * seen with SPI loaders. Today the single entry is {@link IcebergRowLevelDmlTransform}, whose {@code handles} - * is a connector-capability probe (supportsDelete/supportsMerge), not a source-type check.

+ * seen with SPI loaders. Each entry checks the connector's row-change representation and operations, + * not its source name.

*/ public final class RowLevelDmlRegistry { private static final List TRANSFORMS = - ImmutableList.of(new IcebergRowLevelDmlTransform()); + ImmutableList.of(new PositionDeleteRowLevelDmlTransform(), new ChangelogRowLevelDmlTransform()); private RowLevelDmlRegistry() { } @@ -50,6 +54,20 @@ public static Optional find(TableIf table) { return Optional.of(transform); } } + if (table instanceof PluginDrivenExternalTable) { + PluginDrivenExternalTable connectorTable = (PluginDrivenExternalTable) table; + Set operations = connectorTable.connectorSupportedWriteOperations(); + if (supportsAnyRowLevelDml(operations)) { + throw new AnalysisException("No row-level DML plan for connector row-change style " + + connectorTable.getConnectorRowChangeStyle()); + } + } return Optional.empty(); } + + static boolean supportsAnyRowLevelDml(Set operations) { + return operations.contains(WriteOperation.DELETE) + || operations.contains(WriteOperation.UPDATE) + || operations.contains(WriteOperation.MERGE); + } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/RowLevelDmlRowIdUtils.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/RowLevelDmlRowIdUtils.java index 5e169a5f876348..87b4e2d7caf8f8 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/RowLevelDmlRowIdUtils.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/RowLevelDmlRowIdUtils.java @@ -139,7 +139,7 @@ private static boolean pluginConnectorSupportsRowLevelDml(PluginDrivenExternalTa // unresolvable handle (mirroring fetchSyntheticWriteColumns), so a mid-DML catalog drop is "not a target" // rather than an NPE. Set ops = table.connectorSupportedWriteOperations(); - return ops.contains(WriteOperation.DELETE) || ops.contains(WriteOperation.MERGE); + return RowLevelDmlRegistry.supportsAnyRowLevelDml(ops); } /** Check if a plan tree contains any unbound nodes or expressions. */ diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/RowLevelDmlTransform.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/RowLevelDmlTransform.java index 1d7cca65758f3d..841350fd8bebbd 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/RowLevelDmlTransform.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/RowLevelDmlTransform.java @@ -43,7 +43,7 @@ */ public interface RowLevelDmlTransform { - /** Whether this transform handles the given target table (a connector-capability probe). */ + /** Whether this transform handles the table's row-change representation and write operations. */ boolean handles(TableIf table); /** Reject unsupported table modes (e.g. copy-on-write) for the operation, mirroring legacy command checks. */ @@ -59,15 +59,13 @@ BaseExternalTableInsertExecutor newExecutor(ConnectContext ctx, TableIf table, S /** Locate and validate the required physical sink in the planned plan (throws with the legacy messages). */ PhysicalSink requirePhysicalSink(NereidsPlanner planner, RowLevelDmlOp op); - /** The label prefix; the shell appends {@code __}. Frozen for profile/txn parity. */ - String labelPrefix(RowLevelDmlOp op); + /** The connector-owned label prefix; the shell appends {@code __}. */ + String labelPrefix(TableIf table, RowLevelDmlOp op); - /** - * Legacy optimistic-conflict-detection wiring (kept live until P6.7): build the connector-specific - * conflict filter from the analyzed plan and stash it on the executor for its {@code beforeExec}. - */ - void setupConflictDetection(BaseExternalTableInsertExecutor executor, Plan analyzedPlan, TableIf table, - RowLevelDmlOp op); + /** Whether planning must disable external-table batch mode so every source split is available. */ + default boolean requiresExternalTableBatchModeDisabled() { + return false; + } /** Finalize the sink (op-specific; e.g. attaching rewritable delete-file metadata for the BE). */ void finalizeSink(BaseExternalTableInsertExecutor executor, RowLevelDmlOp op, PlanFragment fragment, diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/UpdateCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/UpdateCommand.java index a1b2f01778871c..aa6853cea9f299 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/UpdateCommand.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/UpdateCommand.java @@ -110,7 +110,7 @@ public void run(ConnectContext ctx, StmtExecutor executor) throws Exception { Optional transform = RowLevelDmlRegistry.find(table); if (transform.isPresent()) { RowLevelDmlArgs args = RowLevelDmlArgs.forUpdate( - table, nameParts, tableAlias, assignments, logicalQuery); + table, nameParts, tableAlias, assignments, logicalQuery, cte); new RowLevelDmlCommand(transform.get(), args, RowLevelDmlOp.UPDATE).run(ctx, executor); return; } @@ -280,7 +280,7 @@ public Plan getExplainPlan(ConnectContext ctx) { Optional transform = RowLevelDmlRegistry.find(table); if (transform.isPresent()) { RowLevelDmlArgs args = RowLevelDmlArgs.forUpdate( - table, nameParts, tableAlias, assignments, logicalQuery); + table, nameParts, tableAlias, assignments, logicalQuery, cte); return new RowLevelDmlCommand(transform.get(), args, RowLevelDmlOp.UPDATE).getExplainPlan(ctx); } return completeQueryPlan(ctx, logicalQuery); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/ConnectorChangelogRowChangeSpec.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/ConnectorChangelogRowChangeSpec.java new file mode 100644 index 00000000000000..705db2f57b08b0 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/ConnectorChangelogRowChangeSpec.java @@ -0,0 +1,176 @@ +// 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. + +package org.apache.doris.nereids.trees.plans.commands.info; + +import org.apache.doris.nereids.trees.expressions.EqualTo; +import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.plans.commands.merge.MergeMatchedClause; +import org.apache.doris.nereids.trees.plans.commands.merge.MergeNotMatchedClause; +import org.apache.doris.nereids.util.Utils; + +import com.google.common.collect.ImmutableList; + +import java.util.List; +import java.util.Objects; + +/** Data-only description for a connector whose row-level writes are encoded as changelog rows. */ +public abstract class ConnectorChangelogRowChangeSpec { + public abstract DMLCommandType getDmlCommandType(); + + public abstract List getExpressions(); + + /** UPDATE description. */ + public static final class Update extends ConnectorChangelogRowChangeSpec { + private final List targetNameInPlan; + private final List assignments; + + public Update(List targetNameInPlan, List assignments) { + this.targetNameInPlan = Utils.copyRequiredList(targetNameInPlan); + this.assignments = Utils.copyRequiredList(assignments); + } + + public List getTargetNameInPlan() { + return targetNameInPlan; + } + + public List getAssignments() { + return assignments; + } + + @Override + public DMLCommandType getDmlCommandType() { + return DMLCommandType.UPDATE; + } + + @Override + public List getExpressions() { + return assignments; + } + + @Override + public boolean equals(Object other) { + return other instanceof Update + && Objects.equals(targetNameInPlan, ((Update) other).targetNameInPlan) + && Objects.equals(assignments, ((Update) other).assignments); + } + + @Override + public int hashCode() { + return Objects.hash(targetNameInPlan, assignments); + } + } + + /** DELETE description. */ + public static final class Delete extends ConnectorChangelogRowChangeSpec { + private final List targetNameInPlan; + private final boolean deduplicateTargetRows; + + public Delete(List targetNameInPlan, boolean deduplicateTargetRows) { + this.targetNameInPlan = Utils.copyRequiredList(targetNameInPlan); + this.deduplicateTargetRows = deduplicateTargetRows; + } + + public List getTargetNameInPlan() { + return targetNameInPlan; + } + + public boolean shouldDeduplicateTargetRows() { + return deduplicateTargetRows; + } + + @Override + public DMLCommandType getDmlCommandType() { + return DMLCommandType.DELETE; + } + + @Override + public List getExpressions() { + return ImmutableList.of(); + } + + @Override + public boolean equals(Object other) { + return other instanceof Delete + && Objects.equals(targetNameInPlan, ((Delete) other).targetNameInPlan) + && deduplicateTargetRows == ((Delete) other).deduplicateTargetRows; + } + + @Override + public int hashCode() { + return Objects.hash(targetNameInPlan, deduplicateTargetRows); + } + } + + /** MERGE description. */ + public static final class Merge extends ConnectorChangelogRowChangeSpec { + private final List targetNameInPlan; + private final List matchedClauses; + private final List notMatchedClauses; + + public Merge(List targetNameInPlan, List matchedClauses, + List notMatchedClauses) { + this.targetNameInPlan = Utils.copyRequiredList(targetNameInPlan); + this.matchedClauses = Utils.copyRequiredList(matchedClauses); + this.notMatchedClauses = Utils.copyRequiredList(notMatchedClauses); + } + + public List getTargetNameInPlan() { + return targetNameInPlan; + } + + public List getMatchedClauses() { + return matchedClauses; + } + + public List getNotMatchedClauses() { + return notMatchedClauses; + } + + @Override + public DMLCommandType getDmlCommandType() { + return DMLCommandType.MERGE; + } + + @Override + public List getExpressions() { + ImmutableList.Builder expressions = ImmutableList.builder(); + for (MergeMatchedClause clause : matchedClauses) { + clause.getCasePredicate().ifPresent(expressions::add); + expressions.addAll(clause.getAssignments()); + } + for (MergeNotMatchedClause clause : notMatchedClauses) { + clause.getCasePredicate().ifPresent(expressions::add); + expressions.addAll(clause.getRow()); + } + return expressions.build(); + } + + @Override + public boolean equals(Object other) { + return other instanceof Merge + && Objects.equals(targetNameInPlan, ((Merge) other).targetNameInPlan) + && Objects.equals(matchedClauses, ((Merge) other).matchedClauses) + && Objects.equals(notMatchedClauses, ((Merge) other).notMatchedClauses); + } + + @Override + public int hashCode() { + return Objects.hash(targetNameInPlan, matchedClauses, notMatchedClauses); + } + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/DMLCommandType.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/DMLCommandType.java index aa97f26df18c58..6d9096530071ae 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/DMLCommandType.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/info/DMLCommandType.java @@ -33,6 +33,8 @@ public enum DMLCommandType { UPDATE, // for DELETE DELETE, + // for MERGE INTO + MERGE, // for all other load jobs, including Stream Load, Broker Load, S3 Load // Routine Load etc. LOAD diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/PluginDrivenInsertExecutor.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/PluginDrivenInsertExecutor.java index c85e25512f8f1c..f7df7eb47f4481 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/PluginDrivenInsertExecutor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/PluginDrivenInsertExecutor.java @@ -126,7 +126,7 @@ protected void finalizeSink(PlanFragment fragment, DataSink sink, PhysicalSink p /** * Public finalize entry for the row-level DML shell ({@code RowLevelDmlCommand} via - * {@code IcebergRowLevelDmlTransform.finalizeSink}), which lives outside this package and so cannot reach + * {@code PositionDeleteRowLevelDmlTransform.finalizeSink}), which lives outside this package and so cannot reach * the {@code protected} {@link #finalizeSink}. Mirrors the legacy * {@code IcebergDeleteExecutor.finalizeSinkForDelete} public entry, but with NO rewritable-delete overlay: * the connector's {@code planWrite} supplies {@code rewritable_delete_file_sets} via the write handle (the diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/merge/MergeIntoCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/merge/MergeIntoCommand.java index 24b6adb07760ba..6aa444dbffb314 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/merge/MergeIntoCommand.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/merge/MergeIntoCommand.java @@ -40,7 +40,6 @@ import org.apache.doris.nereids.trees.expressions.NamedExpression; import org.apache.doris.nereids.trees.expressions.Not; import org.apache.doris.nereids.trees.expressions.StatementScopeIdGenerator; -import org.apache.doris.nereids.trees.expressions.functions.scalar.If; import org.apache.doris.nereids.trees.expressions.functions.scalar.Now; import org.apache.doris.nereids.trees.expressions.literal.IntegerLiteral; import org.apache.doris.nereids.trees.expressions.literal.NullLiteral; @@ -212,7 +211,8 @@ private NamedExpression generateBranchLabel(NamedExpression deleteSign) { } Expression currentResult = new IntegerLiteral(i); if (clause.getCasePredicate().isPresent()) { - matchedLabel = new If(clause.getCasePredicate().get(), currentResult, matchedLabel); + matchedLabel = MergeUtils.selectBranch( + clause.getCasePredicate().get(), currentResult, matchedLabel); } else { matchedLabel = currentResult; } @@ -225,12 +225,13 @@ private NamedExpression generateBranchLabel(NamedExpression deleteSign) { } Expression currentResult = new IntegerLiteral(i + matchedClauses.size()); if (clause.getCasePredicate().isPresent()) { - notMatchedLabel = new If(clause.getCasePredicate().get(), currentResult, notMatchedLabel); + notMatchedLabel = MergeUtils.selectBranch( + clause.getCasePredicate().get(), currentResult, notMatchedLabel); } else { notMatchedLabel = currentResult; } } - return new UnboundAlias(new If(new Not(new IsNull(deleteSign)), + return new UnboundAlias(MergeUtils.selectBranch(new Not(new IsNull(deleteSign)), matchedLabel, notMatchedLabel), BRANCH_LABEL); } @@ -447,7 +448,8 @@ private List generateFinalProjections(List colNames, for (int i = 0; i < finalProjections.get(0).size(); i++) { Expression project = new NullLiteral(); for (int j = 0; j < finalProjections.size(); j++) { - project = new If(new EqualTo(new UnboundSlot(BRANCH_LABEL), new IntegerLiteral(j)), + project = MergeUtils.selectBranch( + new EqualTo(new UnboundSlot(BRANCH_LABEL), new IntegerLiteral(j)), finalProjections.get(j).get(i), project); } outputProjectionsBuilder.add(new UnboundAlias(project, colNames.get(i))); diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/merge/MergeUtils.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/merge/MergeUtils.java index 1fbc4c84aea2a5..1fad50aab894d5 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/merge/MergeUtils.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/merge/MergeUtils.java @@ -19,6 +19,7 @@ import org.apache.doris.nereids.rules.exploration.join.JoinReorderContext; import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.functions.scalar.ShortCircuitIf; import org.apache.doris.nereids.trees.plans.JoinType; import org.apache.doris.nereids.trees.plans.logical.LogicalJoin; import org.apache.doris.nereids.trees.plans.logical.LogicalPlan; @@ -35,6 +36,16 @@ public class MergeUtils { private MergeUtils() { } + /** + * Select a MERGE branch without evaluating expressions from branches that were not chosen. + * MERGE clauses are ordered control flow, so their conditions and assignments must not depend + * on the session-wide short_circuit_evaluation setting. + */ + public static Expression selectBranch(Expression condition, + Expression selected, Expression otherwise) { + return new ShortCircuitIf(condition, selected, otherwise); + } + /** * Build the base join between merge target and source, with the target on the LEFT (probe) * side. Doris builds the hash table on the right child, and the target side is structurally diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalConnectorTableSink.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalConnectorTableSink.java index 396d6466beba4b..8961e3be5b3d9f 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalConnectorTableSink.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalConnectorTableSink.java @@ -57,6 +57,7 @@ public class LogicalConnectorTableSink extends LogicalT // can force single-node GATHER output for a rewrite_data_files INSERT-SELECT. Part of plan identity // (equals/hashCode) so the memo never collapses a rewrite sink onto a non-rewrite one. Defaults false. private final boolean rewrite; + private final boolean hasRowOperationColumn; /** * constructor @@ -106,6 +107,25 @@ public LogicalConnectorTableSink(ExternalDatabase database, Optional groupExpression, Optional logicalProperties, CHILD_TYPE child) { + this(database, targetTable, boundTargetSchema, boundPartitionColumns, + boundWriteMetadataIdentity, cols, outputExprs, dmlCommandType, rewrite, + false, groupExpression, logicalProperties, child); + } + + /** Builds a connector sink whose child optionally starts with a row-operation column. */ + public LogicalConnectorTableSink(ExternalDatabase database, + ExternalTable targetTable, + List boundTargetSchema, + List boundPartitionColumns, + String boundWriteMetadataIdentity, + List cols, + List outputExprs, + DMLCommandType dmlCommandType, + boolean rewrite, + boolean hasRowOperationColumn, + Optional groupExpression, + Optional logicalProperties, + CHILD_TYPE child) { super(PlanType.LOGICAL_CONNECTOR_TABLE_SINK, outputExprs, groupExpression, logicalProperties, cols, child); this.database = Objects.requireNonNull(database, "database != null in LogicalConnectorTableSink"); this.targetTable = Objects.requireNonNull(targetTable, "targetTable != null in LogicalConnectorTableSink"); @@ -114,6 +134,7 @@ public LogicalConnectorTableSink(ExternalDatabase database, this.boundWriteMetadataIdentity = boundWriteMetadataIdentity; this.dmlCommandType = dmlCommandType; this.rewrite = rewrite; + this.hasRowOperationColumn = hasRowOperationColumn; } /** Update output expressions based on child output and replace child. */ @@ -124,7 +145,8 @@ public Plan withChildAndUpdateOutput(Plan child) { return AbstractPlan.copyWithSameId(this, () -> new LogicalConnectorTableSink<>(database, targetTable, boundTargetSchema, boundPartitionColumns, boundWriteMetadataIdentity, cols, output, - dmlCommandType, rewrite, Optional.empty(), Optional.empty(), child)); + dmlCommandType, rewrite, hasRowOperationColumn, + Optional.empty(), Optional.empty(), child)); } @Override @@ -133,14 +155,16 @@ public Plan withChildren(List children) { return AbstractPlan.copyWithSameId(this, () -> new LogicalConnectorTableSink<>(database, targetTable, boundTargetSchema, boundPartitionColumns, boundWriteMetadataIdentity, cols, outputExprs, - dmlCommandType, rewrite, Optional.empty(), Optional.empty(), children.get(0))); + dmlCommandType, rewrite, hasRowOperationColumn, + Optional.empty(), Optional.empty(), children.get(0))); } public LogicalConnectorTableSink withOutputExprs(List outputExprs) { return AbstractPlan.copyWithSameId(this, () -> new LogicalConnectorTableSink<>(database, targetTable, boundTargetSchema, boundPartitionColumns, boundWriteMetadataIdentity, cols, outputExprs, - dmlCommandType, rewrite, Optional.empty(), Optional.empty(), child())); + dmlCommandType, rewrite, hasRowOperationColumn, + Optional.empty(), Optional.empty(), child())); } public ExternalDatabase getDatabase() { @@ -171,6 +195,10 @@ public boolean isRewrite() { return rewrite; } + public boolean hasRowOperationColumn() { + return hasRowOperationColumn; + } + @Override public boolean equals(Object o) { if (this == o) { @@ -185,6 +213,7 @@ public boolean equals(Object o) { LogicalConnectorTableSink that = (LogicalConnectorTableSink) o; return dmlCommandType == that.dmlCommandType && rewrite == that.rewrite + && hasRowOperationColumn == that.hasRowOperationColumn && Objects.equals(database, that.database) && Objects.equals(targetTable, that.targetTable) && Objects.equals(boundTargetSchema, that.boundTargetSchema) @@ -196,7 +225,7 @@ public boolean equals(Object o) { @Override public int hashCode() { return Objects.hash(super.hashCode(), database, targetTable, boundTargetSchema, boundPartitionColumns, - boundWriteMetadataIdentity, cols, dmlCommandType, rewrite); + boundWriteMetadataIdentity, cols, dmlCommandType, rewrite, hasRowOperationColumn); } @Override @@ -209,7 +238,8 @@ public String toString() { "boundPartitionColumns", boundPartitionColumns, "cols", cols, "dmlCommandType", dmlCommandType, - "rewrite", rewrite + "rewrite", rewrite, + "hasRowOperationColumn", hasRowOperationColumn ); } @@ -223,7 +253,8 @@ public Plan withGroupExpression(Optional groupExpression) { return AbstractPlan.copyWithSameId(this, () -> new LogicalConnectorTableSink<>(database, targetTable, boundTargetSchema, boundPartitionColumns, boundWriteMetadataIdentity, cols, outputExprs, - dmlCommandType, rewrite, groupExpression, Optional.of(getLogicalProperties()), child())); + dmlCommandType, rewrite, hasRowOperationColumn, + groupExpression, Optional.of(getLogicalProperties()), child())); } @Override @@ -232,6 +263,7 @@ public Plan withGroupExprLogicalPropChildren(Optional groupExpr return AbstractPlan.copyWithSameId(this, () -> new LogicalConnectorTableSink<>(database, targetTable, boundTargetSchema, boundPartitionColumns, boundWriteMetadataIdentity, cols, outputExprs, - dmlCommandType, rewrite, groupExpression, logicalProperties, children.get(0))); + dmlCommandType, rewrite, hasRowOperationColumn, + groupExpression, logicalProperties, children.get(0))); } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalConnectorTableSink.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalConnectorTableSink.java index 9ed1e0d05a233c..8509ae0cb6973c 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalConnectorTableSink.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalConnectorTableSink.java @@ -18,10 +18,14 @@ package org.apache.doris.nereids.trees.plans.physical; import org.apache.doris.catalog.Column; +import org.apache.doris.common.Config; +import org.apache.doris.connector.spi.write.ConnectorWriteDistribution; import org.apache.doris.datasource.ExternalDatabase; import org.apache.doris.datasource.ExternalTable; import org.apache.doris.datasource.plugin.PluginDrivenExternalTable; import org.apache.doris.nereids.memo.GroupExpression; +import org.apache.doris.nereids.properties.DistributionSpecExternalTableSinkHashPartitioned; +import org.apache.doris.nereids.properties.DistributionSpecHash.ShuffleType; import org.apache.doris.nereids.properties.DistributionSpecHiveTableSinkHashPartitioned; import org.apache.doris.nereids.properties.LogicalProperties; import org.apache.doris.nereids.properties.MustLocalSortOrderSpec; @@ -29,19 +33,24 @@ import org.apache.doris.nereids.properties.PhysicalProperties; import org.apache.doris.nereids.trees.expressions.ExprId; import org.apache.doris.nereids.trees.expressions.NamedExpression; +import org.apache.doris.nereids.trees.expressions.Slot; import org.apache.doris.nereids.trees.plans.AbstractPlan; import org.apache.doris.nereids.trees.plans.Plan; import org.apache.doris.nereids.trees.plans.PlanType; +import org.apache.doris.nereids.trees.plans.commands.info.DMLCommandType; import org.apache.doris.nereids.trees.plans.visitor.PlanVisitor; import org.apache.doris.statistics.model.Statistics; +import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; import java.util.ArrayList; import java.util.List; +import java.util.Map; import java.util.Objects; import java.util.Optional; import java.util.Set; +import java.util.TreeMap; import java.util.stream.Collectors; /** @@ -59,6 +68,8 @@ public class PhysicalConnectorTableSink extends Physica // over the partition-shuffle / parallel-write arms below. Carried as a sink field (no ConnectContext, // no instanceof Iceberg). Defaults false → behavior is byte-identical for ordinary connector writes. private final boolean isRewrite; + private final DMLCommandType dmlCommandType; + private final boolean hasRowOperationColumn; /** * constructor @@ -155,12 +166,55 @@ public PhysicalConnectorTableSink(ExternalDatabase database, Statistics statistics, boolean isRewrite, CHILD_TYPE child) { + this(database, targetTable, boundTargetSchema, boundPartitionColumns, + boundWriteMetadataIdentity, cols, outputExprs, groupExpression, logicalProperties, + physicalProperties, statistics, isRewrite, DMLCommandType.NONE, child); + } + + /** Builds a physical connector sink carrying its row-level DML operation. */ + public PhysicalConnectorTableSink(ExternalDatabase database, + ExternalTable targetTable, + List boundTargetSchema, + List boundPartitionColumns, + String boundWriteMetadataIdentity, + List cols, + List outputExprs, + Optional groupExpression, + LogicalProperties logicalProperties, + PhysicalProperties physicalProperties, + Statistics statistics, + boolean isRewrite, + DMLCommandType dmlCommandType, + CHILD_TYPE child) { + this(database, targetTable, boundTargetSchema, boundPartitionColumns, + boundWriteMetadataIdentity, cols, outputExprs, groupExpression, logicalProperties, + physicalProperties, statistics, isRewrite, dmlCommandType, false, child); + } + + /** Builds a physical connector sink with its row shape. */ + public PhysicalConnectorTableSink(ExternalDatabase database, + ExternalTable targetTable, + List boundTargetSchema, + List boundPartitionColumns, + String boundWriteMetadataIdentity, + List cols, + List outputExprs, + Optional groupExpression, + LogicalProperties logicalProperties, + PhysicalProperties physicalProperties, + Statistics statistics, + boolean isRewrite, + DMLCommandType dmlCommandType, + boolean hasRowOperationColumn, + CHILD_TYPE child) { super(PlanType.PHYSICAL_CONNECTOR_TABLE_SINK, database, targetTable, cols, outputExprs, groupExpression, logicalProperties, physicalProperties, statistics, child); this.boundTargetSchema = ImmutableList.copyOf(boundTargetSchema); this.boundPartitionColumns = ImmutableList.copyOf(boundPartitionColumns); this.boundWriteMetadataIdentity = boundWriteMetadataIdentity; this.isRewrite = isRewrite; + this.dmlCommandType = dmlCommandType; + this.hasRowOperationColumn = hasRowOperationColumn; } @Override @@ -169,7 +223,7 @@ public Plan withChildren(List children) { (ExternalDatabase) database, (ExternalTable) targetTable, boundTargetSchema, boundPartitionColumns, boundWriteMetadataIdentity, cols, outputExprs, groupExpression, getLogicalProperties(), physicalProperties, statistics, - isRewrite, children.get(0))); + isRewrite, dmlCommandType, hasRowOperationColumn, children.get(0))); } @Override @@ -182,7 +236,8 @@ public Plan withGroupExpression(Optional groupExpression) { return AbstractPlan.copyWithSameId(this, () -> new PhysicalConnectorTableSink<>( (ExternalDatabase) database, (ExternalTable) targetTable, boundTargetSchema, boundPartitionColumns, boundWriteMetadataIdentity, cols, - outputExprs, groupExpression, getLogicalProperties(), isRewrite, child())); + outputExprs, groupExpression, getLogicalProperties(), PhysicalProperties.GATHER, null, + isRewrite, dmlCommandType, hasRowOperationColumn, child())); } @Override @@ -191,7 +246,8 @@ public Plan withGroupExprLogicalPropChildren(Optional groupExpr return AbstractPlan.copyWithSameId(this, () -> new PhysicalConnectorTableSink<>( (ExternalDatabase) database, (ExternalTable) targetTable, boundTargetSchema, boundPartitionColumns, boundWriteMetadataIdentity, cols, - outputExprs, groupExpression, logicalProperties.get(), isRewrite, children.get(0))); + outputExprs, groupExpression, logicalProperties.get(), PhysicalProperties.GATHER, null, + isRewrite, dmlCommandType, hasRowOperationColumn, children.get(0))); } @Override @@ -200,7 +256,7 @@ public PhysicalPlan withPhysicalPropertiesAndStats(PhysicalProperties physicalPr (ExternalDatabase) database, (ExternalTable) targetTable, boundTargetSchema, boundPartitionColumns, boundWriteMetadataIdentity, cols, outputExprs, groupExpression, getLogicalProperties(), physicalProperties, statistics, - isRewrite, child())); + isRewrite, dmlCommandType, hasRowOperationColumn, child())); } public List getBoundTargetSchema() { @@ -228,6 +284,8 @@ public boolean equals(Object o) { } PhysicalConnectorTableSink that = (PhysicalConnectorTableSink) o; return isRewrite == that.isRewrite + && dmlCommandType == that.dmlCommandType + && hasRowOperationColumn == that.hasRowOperationColumn && Objects.equals(boundTargetSchema, that.boundTargetSchema) && Objects.equals(boundPartitionColumns, that.boundPartitionColumns) && Objects.equals(boundWriteMetadataIdentity, that.boundWriteMetadataIdentity); @@ -236,7 +294,7 @@ public boolean equals(Object o) { @Override public int hashCode() { return Objects.hash(super.hashCode(), boundTargetSchema, boundPartitionColumns, - boundWriteMetadataIdentity, isRewrite); + boundWriteMetadataIdentity, isRewrite, dmlCommandType, hasRowOperationColumn); } /** @@ -248,6 +306,15 @@ public boolean isRewrite() { return isRewrite; } + public DMLCommandType getDmlCommandType() { + return dmlCommandType; + } + + /** Whether the child starts with a connector changelog operation column. */ + public boolean hasRowOperationColumn() { + return hasRowOperationColumn; + } + /** * Get required physical properties for sink distribution. Generalizes the legacy * {@code PhysicalMaxComputeTableSink.getRequirePhysicalProperties()} 3-branch behavior, gated @@ -290,6 +357,9 @@ public PhysicalProperties getRequirePhysicalProperties() { } PluginDrivenExternalTable table = (PluginDrivenExternalTable) targetTable; + Optional connectorDistribution + = table.getConnectorWriteDistribution(); + if (table.requirePartitionLocalSortOnWrite()) { Set partitionNames = boundPartitionColumns.stream() .map(Column::getName) @@ -318,25 +388,36 @@ public PhysicalProperties getRequirePhysicalProperties() { } } List exprIds = columnIdx.stream() - .map(idx -> child().getOutput().get(idx).getExprId()) + .map(idx -> child().getOutput().get( + idx + (hasRowOperationColumn() ? 1 : 0)).getExprId()) .collect(Collectors.toList()); - DistributionSpecHiveTableSinkHashPartitioned shuffleInfo - = new DistributionSpecHiveTableSinkHashPartitioned(); - shuffleInfo.setOutputColExprIds(exprIds); + PhysicalProperties requiredProperties; + if (connectorDistribution.isPresent()) { + requiredProperties = toPhysicalProperties(connectorDistribution.get()); + } else { + DistributionSpecHiveTableSinkHashPartitioned shuffleInfo + = new DistributionSpecHiveTableSinkHashPartitioned(); + shuffleInfo.setOutputColExprIds(exprIds); + requiredProperties = new PhysicalProperties(shuffleInfo); + } // Local sort by partition columns so rows for the same partition are grouped // together before the streaming partition writer (MaxCompute Storage API closes a // partition writer once a different partition value appears). List orderKeys = columnIdx.stream() - .map(idx -> new OrderKey(child().getOutput().get(idx), true, false)) + .map(idx -> new OrderKey(child().getOutput().get( + idx + (hasRowOperationColumn() ? 1 : 0)), true, false)) .collect(Collectors.toList()); - return new PhysicalProperties(shuffleInfo) - .withOrderSpec(new MustLocalSortOrderSpec(orderKeys)); + return requiredProperties.withOrderSpec(new MustLocalSortOrderSpec(orderKeys)); } // Partition columns exist but none in cols == all partitions statically specified; // fall through to the parallel/gather branch (no sort/shuffle needed). } } + if (connectorDistribution.isPresent()) { + return toPhysicalProperties(connectorDistribution.get()); + } + if (table.requirePartitionHashOnWrite()) { Set partitionNames = boundPartitionColumns.stream() .map(Column::getName) @@ -356,7 +437,8 @@ public PhysicalProperties getRequirePhysicalProperties() { } } List exprIds = columnIdx.stream() - .map(idx -> child().getOutput().get(idx).getExprId()) + .map(idx -> child().getOutput().get( + idx + (hasRowOperationColumn() ? 1 : 0)).getExprId()) .collect(Collectors.toList()); DistributionSpecHiveTableSinkHashPartitioned shuffleInfo = new DistributionSpecHiveTableSinkHashPartitioned(); @@ -370,4 +452,56 @@ public PhysicalProperties getRequirePhysicalProperties() { } return PhysicalProperties.GATHER; } + + private PhysicalProperties toPhysicalProperties(ConnectorWriteDistribution distribution) { + switch (distribution.getMode()) { + case EXECUTION_ANY: + return PhysicalProperties.EXECUTION_ANY; + case GATHER: + return PhysicalProperties.GATHER; + case HASH: + return PhysicalProperties.createHash( + routeExprIds(distribution.getRouteColumns()), ShuffleType.REQUIRE); + case EXTERNAL_UNPARTITIONED: + requireExternalWriterRoutingSupport(); + return PhysicalProperties.EXTERNAL_TABLE_SINK_UNPARTITIONED; + case EXTERNAL_HASH: + requireExternalWriterRoutingSupport(); + return new PhysicalProperties(new DistributionSpecExternalTableSinkHashPartitioned( + routeExprIds(distribution.getRouteColumns()), + distribution.getPartitionFunction(), + distribution.getPartitionFunctionOptions(), + distribution.getWriterAssignment())); + default: + throw new IllegalStateException("Unsupported connector write distribution: " + + distribution.getMode()); + } + } + + private void requireExternalWriterRoutingSupport() { + Preconditions.checkState(Config.be_exec_version + >= DistributionSpecExternalTableSinkHashPartitioned.MIN_BE_EXEC_VERSION, + "External table sink distribution requires BE execution version %s or newer", + DistributionSpecExternalTableSinkHashPartitioned.MIN_BE_EXEC_VERSION); + } + + private List routeExprIds(List routeColumns) { + List output = child().getOutput(); + int offset = hasRowOperationColumn() ? 1 : 0; + PluginDrivenExternalTable table = (PluginDrivenExternalTable) targetTable; + List outputColumns = table.requiresFullSchemaWriteOrder() + ? boundTargetSchema : cols; + Preconditions.checkState(outputColumns.size() + offset == output.size(), + "Connector sink schema must match child output for routed writes"); + Map outputByName = new TreeMap<>(String.CASE_INSENSITIVE_ORDER); + for (int i = 0; i < outputColumns.size(); i++) { + outputByName.put(outputColumns.get(i).getName(), output.get(i + offset).getExprId()); + } + List exprIds = new ArrayList<>(routeColumns.size()); + for (String column : routeColumns) { + exprIds.add(Preconditions.checkNotNull(outputByName.get(column), + "Connector route column is missing from sink output: " + column)); + } + return exprIds; + } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/types/VariantType.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/types/VariantType.java index 37807269ed85c7..346435bcb1827b 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/types/VariantType.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/types/VariantType.java @@ -305,4 +305,30 @@ public boolean getEnableNestedGroup() { return enableNestedGroup; } + /** Whether the Variant V2 execution kernel can convert this source type. */ + public static boolean isSupportedComputeV2CastSource(DataType dataType) { + if (dataType.isNullType() || dataType.isJsonType()) { + return true; + } + if (dataType instanceof VariantType) { + // Master is V2-only. Ordinary Variant and the connector execution marker share the + // same runtime representation, regardless of storage-layout properties. + return true; + } + if (dataType instanceof ArrayType) { + return isSupportedComputeV2CastSource(((ArrayType) dataType).getItemType()); + } + if (dataType instanceof DecimalV3Type) { + return ((DecimalV3Type) dataType).getPrecision() + <= DecimalV3Type.MAX_DECIMAL128_PRECISION; + } + return dataType.isBooleanType() + || dataType.isIntegralType() + || dataType.isFloatLikeType() + || dataType.isDecimalV2Type() + || dataType.isDateLikeType() + || dataType.isStringLikeType() + || dataType.isIPType(); + } + } diff --git a/fe/fe-core/src/main/java/org/apache/doris/planner/DataPartition.java b/fe/fe-core/src/main/java/org/apache/doris/planner/DataPartition.java index 0ef85f8ee67170..9a2385b8a3350f 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/planner/DataPartition.java +++ b/fe/fe-core/src/main/java/org/apache/doris/planner/DataPartition.java @@ -26,6 +26,8 @@ import org.apache.doris.analysis.ToSqlParams; import org.apache.doris.thrift.TDataPartition; import org.apache.doris.thrift.TExplainLevel; +import org.apache.doris.thrift.TExternalTableSinkHashPartitionInfo; +import org.apache.doris.thrift.TExternalTableSinkWriterAssignment; import org.apache.doris.thrift.TIcebergPartitionField; import org.apache.doris.thrift.TMergePartitionInfo; import org.apache.doris.thrift.TPartitionType; @@ -33,9 +35,11 @@ import com.google.common.base.Joiner; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; import com.google.common.collect.Lists; import java.util.List; +import java.util.Map; /** * Specification of the partition of a single stream of data. @@ -55,6 +59,9 @@ public class DataPartition { // for hash partition: exprs used to compute hash value private ImmutableList partitionExprs; private MergePartitionInfo mergePartitionInfo; + private String externalPartitionFunction; + private ImmutableMap externalPartitionFunctionOptions = ImmutableMap.of(); + private TExternalTableSinkWriterAssignment externalWriterAssignment; public DataPartition(TPartitionType type, List exprs) { Preconditions.checkNotNull(exprs); @@ -62,15 +69,27 @@ public DataPartition(TPartitionType type, List exprs) { Preconditions.checkState(type == TPartitionType.HASH_PARTITIONED || type == TPartitionType.RANGE_PARTITIONED || type == TPartitionType.HIVE_TABLE_SINK_HASH_PARTITIONED + || type == TPartitionType.EXTERNAL_TABLE_SINK_HASH_PARTITIONED || type == TPartitionType.BUCKET_SHFFULE_HASH_PARTITIONED); this.type = type; this.partitionExprs = ImmutableList.copyOf(exprs); } + public DataPartition(TPartitionType type, List exprs, String partitionFunction, + Map partitionFunctionOptions, + TExternalTableSinkWriterAssignment writerAssignment) { + this(type, exprs); + Preconditions.checkState(type == TPartitionType.EXTERNAL_TABLE_SINK_HASH_PARTITIONED); + this.externalPartitionFunction = Preconditions.checkNotNull(partitionFunction); + this.externalPartitionFunctionOptions = ImmutableMap.copyOf(partitionFunctionOptions); + this.externalWriterAssignment = Preconditions.checkNotNull(writerAssignment); + } + public DataPartition(TPartitionType type) { Preconditions.checkState(type == TPartitionType.UNPARTITIONED || type == TPartitionType.RANDOM || type == TPartitionType.HIVE_TABLE_SINK_UNPARTITIONED + || type == TPartitionType.EXTERNAL_TABLE_SINK_UNPARTITIONED || type == TPartitionType.OLAP_TABLE_SINK_HASH_PARTITIONED); this.type = type; this.partitionExprs = ImmutableList.of(); @@ -110,6 +129,13 @@ public TDataPartition toThrift() { if (mergePartitionInfo != null) { result.setMergePartitionInfo(mergePartitionInfo.toThrift()); } + if (externalPartitionFunction != null) { + TExternalTableSinkHashPartitionInfo info = new TExternalTableSinkHashPartitionInfo(); + info.setPartitionFunction(externalPartitionFunction); + info.setPartitionFunctionOptions(externalPartitionFunctionOptions); + info.setWriterAssignment(externalWriterAssignment); + result.setExternalTableSinkHashPartitionInfo(info); + } return result; } diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/AbstractJobProcessor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/AbstractJobProcessor.java index 647391dffbc422..0231704d03021f 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/AbstractJobProcessor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/AbstractJobProcessor.java @@ -28,6 +28,7 @@ import org.apache.doris.thrift.TStatus; import org.apache.doris.thrift.TStatusCode; import org.apache.doris.thrift.TUniqueId; +import org.apache.doris.transaction.CommitDataSerializer; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableMap; @@ -108,8 +109,7 @@ public final boolean updateFragmentExecStatus(TReportExecStatusParams params) { SingleFragmentPipelineTask fragmentTask = backendFragmentTasks.get().get( new BackendFragmentId(params.getBackendId(), params.getFragmentId())); if (fragmentTask == null) { - if (params.isSetHivePartitionUpdates() || params.isSetIcebergCommitDatas() - || params.isSetMcCommitDatas()) { + if (CommitDataSerializer.hasCommitData(params)) { throw new IllegalStateException("Missing fragment handler for external-file report"); } return false; @@ -138,8 +138,7 @@ public final boolean updateFragmentExecStatus(TReportExecStatusParams params) { } } doProcessReportExecStatus(params, fragmentTask); - return !params.isSetHivePartitionUpdates() && !params.isSetIcebergCommitDatas() - && !params.isSetMcCommitDatas() || fragmentTask.isDone(); + return !CommitDataSerializer.hasCommitData(params) || fragmentTask.isDone(); } private Map buildBackendFragmentTasks( diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java b/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java index 0acfe3d7635087..cc291c23881c07 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/Coordinator.java @@ -2771,8 +2771,7 @@ public boolean updateFragmentExecStatus(TReportExecStatusParams params) { } PipelineExecContext ctx = pipelineExecContexts.get(Pair.of(params.getFragmentId(), params.getBackendId())); - boolean hasExternalCommitData = params.isSetHivePartitionUpdates() - || params.isSetIcebergCommitDatas() || params.isSetMcCommitDatas(); + boolean hasExternalCommitData = CommitDataSerializer.hasCommitData(params); if (ctx == null) { if (hasExternalCommitData) { throw new IllegalStateException("Missing fragment handler for external-file report"); @@ -2842,18 +2841,9 @@ public boolean updateFragmentExecStatus(TReportExecStatusParams params) { if (params.isSetErrorTabletInfos()) { updateErrorTabletInfos(params.getErrorTabletInfos()); } - if (params.isSetHivePartitionUpdates() || params.isSetIcebergCommitDatas() - || params.isSetMcCommitDatas()) { + if (CommitDataSerializer.hasCommitData(params)) { Transaction txn = Env.getCurrentEnv().getGlobalExternalTransactionInfoMgr().getTxnById(reportTxnId); - if (params.isSetHivePartitionUpdates()) { - CommitDataSerializer.feed(txn, params.getHivePartitionUpdates()); - } - if (params.isSetIcebergCommitDatas()) { - CommitDataSerializer.feed(txn, params.getIcebergCommitDatas()); - } - if (params.isSetMcCommitDatas()) { - CommitDataSerializer.feed(txn, params.getMcCommitDatas()); - } + CommitDataSerializer.feed(txn, params); } accepted = true; diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/QeProcessorImpl.java b/fe/fe-core/src/main/java/org/apache/doris/qe/QeProcessorImpl.java index 52bfcf4a0ff545..77a4e0b4538657 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/QeProcessorImpl.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/QeProcessorImpl.java @@ -35,6 +35,7 @@ import org.apache.doris.thrift.TStatus; import org.apache.doris.thrift.TStatusCode; import org.apache.doris.thrift.TUniqueId; +import org.apache.doris.transaction.CommitDataSerializer; import com.google.common.base.Strings; import com.google.common.cache.Cache; @@ -337,7 +338,7 @@ public TReportExecStatusResult reportExecStatus(TReportExecStatusParams params, } private static boolean hasExternalCommitData(TReportExecStatusParams params) { - return params.isSetHivePartitionUpdates() || params.isSetIcebergCommitDatas() || params.isSetMcCommitDatas(); + return CommitDataSerializer.hasCommitData(params); } private static String externalFileReportKey(TReportExecStatusParams params) { diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/LoadProcessor.java b/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/LoadProcessor.java index d4878ff99a3d6c..d22ddb7689e64b 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/LoadProcessor.java +++ b/fe/fe-core/src/main/java/org/apache/doris/qe/runtime/LoadProcessor.java @@ -187,8 +187,7 @@ protected void doProcessReportExecStatus(TReportExecStatusParams params, SingleF } if (!fragmentTask.processReportExecStatus(params, () -> acceptFinalReport(params))) { - if ((params.isSetHivePartitionUpdates() || params.isSetIcebergCommitDatas() - || params.isSetMcCommitDatas()) && !fragmentTask.isDone()) { + if (CommitDataSerializer.hasCommitData(params) && !fragmentTask.isDone()) { throw new IllegalStateException("External-file report was not a completed fragment report"); } LOG.debug("Fragment {} is not done, ignore report status: {}", @@ -246,17 +245,9 @@ private void acceptFinalReport(TReportExecStatusParams params) { loadContext.updateErrorTabletInfos(params.getErrorTabletInfos()); } long txnId = loadContext.getTransactionId(); - if (params.isSetHivePartitionUpdates() || params.isSetIcebergCommitDatas() || params.isSetMcCommitDatas()) { + if (CommitDataSerializer.hasCommitData(params)) { Transaction txn = Env.getCurrentEnv().getGlobalExternalTransactionInfoMgr().getTxnById(txnId); - if (params.isSetHivePartitionUpdates()) { - CommitDataSerializer.feed(txn, params.getHivePartitionUpdates()); - } - if (params.isSetIcebergCommitDatas()) { - CommitDataSerializer.feed(txn, params.getIcebergCommitDatas()); - } - if (params.isSetMcCommitDatas()) { - CommitDataSerializer.feed(txn, params.getMcCommitDatas()); - } + CommitDataSerializer.feed(txn, params); } } diff --git a/fe/fe-core/src/main/java/org/apache/doris/transaction/CommitDataSerializer.java b/fe/fe-core/src/main/java/org/apache/doris/transaction/CommitDataSerializer.java index 1e4bc17295aa1c..a9738ec030a720 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/transaction/CommitDataSerializer.java +++ b/fe/fe-core/src/main/java/org/apache/doris/transaction/CommitDataSerializer.java @@ -17,11 +17,14 @@ package org.apache.doris.transaction; +import org.apache.doris.thrift.TReportExecStatusParams; + import org.apache.thrift.TBase; import org.apache.thrift.TException; import org.apache.thrift.TSerializer; import org.apache.thrift.protocol.TBinaryProtocol; +import java.nio.ByteBuffer; import java.util.List; import java.util.stream.Collectors; @@ -40,6 +43,28 @@ public final class CommitDataSerializer { private CommitDataSerializer() { } + /** Returns whether a fragment report carries any external connector commit data. */ + public static boolean hasCommitData(TReportExecStatusParams params) { + return params.isSetHivePartitionUpdates() || params.isSetIcebergCommitDatas() + || params.isSetMcCommitDatas() || params.isSetConnectorCommitData(); + } + + /** Delivers every commit-data representation carried by one fragment report. */ + public static void feed(Transaction txn, TReportExecStatusParams params) { + if (params.isSetHivePartitionUpdates()) { + feed(txn, params.getHivePartitionUpdates()); + } + if (params.isSetIcebergCommitDatas()) { + feed(txn, params.getIcebergCommitDatas()); + } + if (params.isSetMcCommitDatas()) { + feed(txn, params.getMcCommitDatas()); + } + if (params.isSetConnectorCommitData()) { + feedRaw(txn, params.getConnectorCommitData()); + } + } + /** * Serializes each commit fragment and accumulates it into {@code txn}. * @@ -67,6 +92,20 @@ public static void feed(Transaction txn, List> fragments) } } + /** + * Delivers opaque commit fragments without interpreting connector-owned bytes in FE core. + * Thrift exposes binary values as {@link ByteBuffer}; copy each remaining slice before + * passing it to a transaction, which may keep the byte array after the RPC is released. + */ + public static void feedRaw(Transaction txn, List fragments) { + for (ByteBuffer fragment : fragments) { + ByteBuffer source = fragment.duplicate(); + byte[] bytes = new byte[source.remaining()]; + source.get(bytes); + txn.addCommitData(bytes); + } + } + private static final class CommitDataSerializationException extends RuntimeException { private CommitDataSerializationException(TException cause) { super(cause); diff --git a/fe/fe-core/src/main/java/org/apache/doris/transaction/PluginDrivenTransactionManager.java b/fe/fe-core/src/main/java/org/apache/doris/transaction/PluginDrivenTransactionManager.java index f5e81a6ebb7067..69a7b745fde75b 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/transaction/PluginDrivenTransactionManager.java +++ b/fe/fe-core/src/main/java/org/apache/doris/transaction/PluginDrivenTransactionManager.java @@ -138,6 +138,7 @@ public boolean isActive(long id) { private static class PluginDrivenTransaction implements Transaction { private final long id; protected final ConnectorTransaction connectorTx; + private boolean active = true; PluginDrivenTransaction(long id, ConnectorTransaction connectorTx) { this.id = id; @@ -145,7 +146,9 @@ private static class PluginDrivenTransaction implements Transaction { } @Override - public void commit() { + public synchronized void commit() { + requireActive(); + active = false; if (connectorTx == null) { return; } @@ -157,7 +160,9 @@ public void commit() { } @Override - public void rollback() { + public synchronized void rollback() { + requireActive(); + active = false; if (connectorTx == null) { return; } @@ -169,7 +174,8 @@ public void rollback() { } @Override - public void addCommitData(byte[] commitFragment) { + public synchronized void addCommitData(byte[] commitFragment) { + requireActive(); if (connectorTx != null) { connectorTx.addCommitData(commitFragment); } @@ -177,10 +183,16 @@ public void addCommitData(byte[] commitFragment) { } @Override - public long getUpdateCnt() { + public synchronized long getUpdateCnt() { return connectorTx == null ? 0 : connectorTx.getUpdateCnt(); } + protected synchronized void requireActive() { + if (!active) { + throw new IllegalStateException("Plugin-driven transaction is already finished: " + id); + } + } + private void closeQuietly() { try { connectorTx.close(); @@ -205,7 +217,8 @@ private static final class WriteBlockAllocatingPluginDrivenTransaction extends P } @Override - public long allocateWriteBlockRange(String writeSessionId, long count) throws UserException { + public synchronized long allocateWriteBlockRange(String writeSessionId, long count) throws UserException { + requireActive(); return ((WriteBlockAllocatingConnectorTransaction) connectorTx) .allocateWriteBlockRange(writeSessionId, count); } diff --git a/fe/fe-core/src/test/java/org/apache/doris/datasource/plugin/PluginDrivenExternalTableTest.java b/fe/fe-core/src/test/java/org/apache/doris/datasource/plugin/PluginDrivenExternalTableTest.java index 15e52ef34b9554..ad8a0f1f65243a 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/datasource/plugin/PluginDrivenExternalTableTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/datasource/plugin/PluginDrivenExternalTableTest.java @@ -35,6 +35,7 @@ import org.apache.doris.connector.spi.mvcc.ConnectorMvccSnapshot; import org.apache.doris.connector.spi.pushdown.ConnectorColumnRef; import org.apache.doris.connector.spi.pushdown.ConnectorExpression; +import org.apache.doris.connector.spi.write.ConnectorRowChangeStyle; import org.apache.doris.connector.spi.write.ConnectorWritePlanProvider; import org.apache.doris.datasource.ExternalDatabase; import org.apache.doris.datasource.SchemaCacheValue; @@ -107,6 +108,11 @@ private static ConnectorSession noneScopedSession() { */ private static PluginDrivenExternalTable capabilityTable(boolean handlePresent, Set ops, boolean branch) { + return capabilityTable(handlePresent, ops, branch, ConnectorRowChangeStyle.POSITION_DELETE); + } + + private static PluginDrivenExternalTable capabilityTable(boolean handlePresent, + Set ops, boolean branch, ConnectorRowChangeStyle style) { ConnectorTableHandle handle = Mockito.mock(ConnectorTableHandle.class); ConnectorMetadata metadata = Mockito.mock(ConnectorMetadata.class); Mockito.when(metadata.getTableHandle(Mockito.any(), Mockito.any(), Mockito.any())) @@ -115,6 +121,7 @@ private static PluginDrivenExternalTable capabilityTable(boolean handlePresent, // provider, so stub them where they are actually declared. ConnectorWritePlanProvider provider = Mockito.mock(ConnectorWritePlanProvider.class); Mockito.when(provider.supportedOperations()).thenReturn(ops); + Mockito.when(provider.getRowChangeStyle()).thenReturn(style); Mockito.when(provider.supportsWriteBranch()).thenReturn(branch); Mockito.when(provider.requiresPartitionHashWrite()).thenReturn(true); Mockito.when(provider.requiresMaterializeStaticPartitionValues()).thenReturn(true); @@ -137,6 +144,7 @@ public void connectorWriteCapabilitiesResolvePerHandle() { PluginDrivenExternalTable table = capabilityTable(true, ops, true); Assertions.assertEquals(ops, table.connectorSupportedWriteOperations(), "the write ops must come from the connector's per-handle overload (resolved via the handle)"); + Assertions.assertEquals(ConnectorRowChangeStyle.POSITION_DELETE, table.getConnectorRowChangeStyle()); Assertions.assertTrue(table.connectorSupportsWriteBranch(), "the branch capability must come from the connector's per-handle overload"); Assertions.assertTrue(table.requirePartitionHashOnWrite(), @@ -152,6 +160,7 @@ public void connectorWriteCapabilitiesDegradeWhenHandleUnresolvable() { PluginDrivenExternalTable table = capabilityTable(false, EnumSet.of(WriteOperation.DELETE), true); Assertions.assertTrue(table.connectorSupportedWriteOperations().isEmpty(), "an unresolvable handle degrades write ops to the empty set"); + Assertions.assertEquals(ConnectorRowChangeStyle.NONE, table.getConnectorRowChangeStyle()); Assertions.assertFalse(table.connectorSupportsWriteBranch(), "an unresolvable handle degrades branch support to false"); Assertions.assertFalse(table.requirePartitionHashOnWrite(), @@ -170,9 +179,19 @@ public void connectorWriteCapabilitiesDegradeWhenConnectorNull() { Deencapsulation.setField(table, "catalog", catalog); Assertions.assertTrue(table.connectorSupportedWriteOperations().isEmpty(), "a null connector degrades write ops to the empty set"); + Assertions.assertEquals(ConnectorRowChangeStyle.NONE, table.getConnectorRowChangeStyle()); Assertions.assertFalse(table.connectorSupportsWriteBranch(), "a null connector degrades branch to false"); } + @Test + public void rowChangeStyleComesFromThePerHandleWriteProvider() { + PluginDrivenExternalTable table = capabilityTable(true, + EnumSet.of(WriteOperation.INSERT, WriteOperation.DELETE), false, + ConnectorRowChangeStyle.CHANGELOG); + + Assertions.assertEquals(ConnectorRowChangeStyle.CHANGELOG, table.getConnectorRowChangeStyle()); + } + // ==================== §4.4 W4: per-handle transaction write-target handle resolution ==================== // A CALLS_REAL_METHODS table whose connector resolves the write-target handle to `resolved` (null => empty). @@ -422,7 +441,11 @@ public void resolveWriteColumnsRunsAndRestoresPluginContextClassLoader() { ConnectorSession session = noneScopedSession(); Connector connector = Mockito.mock(Connector.class); Mockito.when(connector.getMetadata(Mockito.any())).thenReturn(metadata); - Mockito.when(connector.getWritePlanProvider(handle)).thenReturn(provider); + AtomicReference providerResolutionLoader = new AtomicReference<>(); + Mockito.when(connector.getWritePlanProvider(handle)).thenAnswer(invocation -> { + providerResolutionLoader.set(Thread.currentThread().getContextClassLoader()); + return provider; + }); PluginDrivenExternalCatalog catalog = Mockito.mock(PluginDrivenExternalCatalog.class); Mockito.when(catalog.getConnector()).thenReturn(connector); Mockito.when(catalog.buildConnectorSession()).thenReturn(session); @@ -439,6 +462,7 @@ public void resolveWriteColumnsRunsAndRestoresPluginContextClassLoader() { Thread.currentThread().setContextClassLoader(previous); try { Assertions.assertTrue(table.resolveWriteColumns(Optional.empty()).isPresent()); + Assertions.assertSame(connector.getClass().getClassLoader(), providerResolutionLoader.get()); Assertions.assertSame(provider.getClass().getClassLoader(), observed.get()); Assertions.assertEquals("pinned-generation", ctx.getStatementContext() .getConnectorWriteMetadataIdentity(99L).orElse(null)); diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/glue/translator/ExpressionTranslatorTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/glue/translator/ExpressionTranslatorTest.java index 06236be1f47967..1b66750aa4a3b4 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/glue/translator/ExpressionTranslatorTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/glue/translator/ExpressionTranslatorTest.java @@ -20,7 +20,10 @@ import org.apache.doris.analysis.ArithmeticExpr; import org.apache.doris.analysis.ArithmeticExpr.Operator; import org.apache.doris.analysis.Expr; +import org.apache.doris.analysis.ExprToThriftVisitor; +import org.apache.doris.analysis.FunctionCallExpr; import org.apache.doris.analysis.IntLiteral; +import org.apache.doris.analysis.ShortCircuitFunctionCallExpr; import org.apache.doris.analysis.SlotRef; import org.apache.doris.catalog.Function.NullableMode; import org.apache.doris.catalog.Type; @@ -29,10 +32,13 @@ import org.apache.doris.nereids.trees.expressions.MatchAny; import org.apache.doris.nereids.trees.expressions.Or; import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.functions.scalar.ShortCircuitIf; +import org.apache.doris.nereids.trees.expressions.literal.BooleanLiteral; import org.apache.doris.nereids.trees.expressions.literal.IntegerLiteral; import org.apache.doris.nereids.trees.expressions.literal.NullLiteral; import org.apache.doris.nereids.trees.expressions.literal.VarcharLiteral; import org.apache.doris.nereids.types.IntegerType; +import org.apache.doris.thrift.TExprNode; import com.google.common.collect.ImmutableList; import org.junit.jupiter.api.Assertions; @@ -72,4 +78,18 @@ public void testMatch() { Expr actual = translator.visitOr(or, context); Assertions.assertTrue(actual.isNullable()); } + + @Test + void testRequiredShortCircuitEvaluationSurvivesTranslation() { + ShortCircuitIf expression = new ShortCircuitIf( + BooleanLiteral.TRUE, new IntegerLiteral(1), new IntegerLiteral(2)); + + FunctionCallExpr translated = (FunctionCallExpr) ExpressionTranslator.translate( + expression, new PlanTranslatorContext()); + TExprNode thriftNode = ExprToThriftVisitor.treeToThrift(translated).getNodes().get(0); + + Assertions.assertInstanceOf(ShortCircuitFunctionCallExpr.class, translated); + Assertions.assertTrue(thriftNode.isSetShortCircuitEvaluation()); + Assertions.assertTrue(thriftNode.isShortCircuitEvaluation()); + } } diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/analysis/ConnectorChangelogPlanBuilderTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/analysis/ConnectorChangelogPlanBuilderTest.java new file mode 100644 index 00000000000000..4fa2f7fcacec3f --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/analysis/ConnectorChangelogPlanBuilderTest.java @@ -0,0 +1,94 @@ +// 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. + +package org.apache.doris.nereids.rules.analysis; + +import org.apache.doris.catalog.Column; +import org.apache.doris.catalog.PrimitiveType; +import org.apache.doris.catalog.ScalarType; +import org.apache.doris.connector.spi.write.ConnectorChangelogMode; +import org.apache.doris.nereids.CascadesContext; +import org.apache.doris.nereids.analyzer.UnboundSlot; +import org.apache.doris.nereids.trees.expressions.EqualTo; +import org.apache.doris.nereids.trees.expressions.NamedExpression; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.literal.IntegerLiteral; +import org.apache.doris.nereids.trees.plans.RelationId; +import org.apache.doris.nereids.trees.plans.commands.info.ConnectorChangelogRowChangeSpec; +import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate; +import org.apache.doris.nereids.trees.plans.logical.LogicalEmptyRelation; +import org.apache.doris.nereids.trees.plans.logical.LogicalPlan; +import org.apache.doris.nereids.trees.plans.logical.LogicalProject; +import org.apache.doris.nereids.types.IntegerType; +import org.apache.doris.nereids.util.MemoTestUtils; + +import com.google.common.collect.ImmutableList; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.List; + +class ConnectorChangelogPlanBuilderTest { + private static final ConnectorChangelogMode MODE = + new ConnectorChangelogMode("connector_operation", (byte) 4, (byte) 6, (byte) 8); + private static final List SCHEMA = ImmutableList.of( + new Column("id", ScalarType.createType(PrimitiveType.INT)), + new Column("value", ScalarType.createType(PrimitiveType.INT))); + + @Test + void updateUsesConnectorOwnedOperationEncoding() { + LogicalPlan child = targetRow(); + CascadesContext context = MemoTestUtils.createCascadesContext(child); + ConnectorChangelogRowChangeSpec.Update spec = new ConnectorChangelogRowChangeSpec.Update( + ImmutableList.of("target"), ImmutableList.of( + new EqualTo(new UnboundSlot("value"), new IntegerLiteral(99)))); + + LogicalPlan result = ConnectorChangelogPlanBuilder.build( + SCHEMA, ImmutableList.of("id"), MODE, spec, child, context); + + Assertions.assertInstanceOf(LogicalProject.class, result); + Assertions.assertEquals(ImmutableList.of("connector_operation", "id", "value"), + result.getOutput().stream().map(NamedExpression::getName) + .collect(ImmutableList.toImmutableList())); + Assertions.assertTrue(((LogicalProject) result).getProjects().get(0).toSql().contains("6")); + Assertions.assertTrue(((LogicalProject) result).getProjects().get(2).toSql().contains("99")); + } + + @Test + void deleteUsingDeduplicatesByConnectorPrimaryKey() { + LogicalPlan child = targetRow(); + CascadesContext context = MemoTestUtils.createCascadesContext(child); + ConnectorChangelogRowChangeSpec.Delete spec = new ConnectorChangelogRowChangeSpec.Delete( + ImmutableList.of("target"), true); + + LogicalPlan result = ConnectorChangelogPlanBuilder.build( + SCHEMA, ImmutableList.of("id"), MODE, spec, child, context); + + Assertions.assertInstanceOf(LogicalAggregate.class, result); + Assertions.assertEquals(ImmutableList.of("connector_operation", "id", "value"), + result.getOutput().stream().map(NamedExpression::getName) + .collect(ImmutableList.toImmutableList())); + } + + private LogicalPlan targetRow() { + SlotReference id = new SlotReference("id", IntegerType.INSTANCE, false, + ImmutableList.of("target")); + SlotReference value = new SlotReference("value", IntegerType.INSTANCE, true, + ImmutableList.of("target")); + return new LogicalEmptyRelation(new RelationId(1), ImmutableList.of(id, value)); + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/analysis/FunctionRegistryTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/analysis/FunctionRegistryTest.java index 6c8eaaac0769fd..df738f9966d90d 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/analysis/FunctionRegistryTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/analysis/FunctionRegistryTest.java @@ -37,8 +37,11 @@ import org.apache.doris.nereids.trees.expressions.functions.scalar.Year; import org.apache.doris.nereids.trees.expressions.literal.Literal; import org.apache.doris.nereids.trees.expressions.shape.UnaryExpression; +import org.apache.doris.nereids.types.ArrayType; import org.apache.doris.nereids.types.BitmapType; import org.apache.doris.nereids.types.IntegerType; +import org.apache.doris.nereids.types.MapType; +import org.apache.doris.nereids.types.StructType; import org.apache.doris.nereids.util.MemoPatternMatchSupported; import org.apache.doris.nereids.util.MemoTestUtils; import org.apache.doris.nereids.util.PlanChecker; @@ -148,6 +151,39 @@ public void testVariantParseFunctions() { ); } + @Test + public void testVariantNestedConstructors() { + PlanChecker.from(connectContext) + .analyze("select array(parse_to_variant('1')), " + + "map('k', parse_to_variant('2')), " + + "struct(parse_to_variant('3')), " + + "named_struct('v', parse_to_variant('4'))") + .matches( + logicalOneRowRelation().when(oneRowRelation -> { + ArrayType array = (ArrayType) oneRowRelation.getProjects().get(0) + .child(0).getDataType(); + MapType map = (MapType) oneRowRelation.getProjects().get(1) + .child(0).getDataType(); + StructType struct = (StructType) oneRowRelation.getProjects().get(2) + .child(0).getDataType(); + StructType namedStruct = (StructType) oneRowRelation.getProjects().get(3) + .child(0).getDataType(); + Assertions.assertTrue(array.getItemType().isVariantType()); + Assertions.assertTrue(map.getValueType().isVariantType()); + Assertions.assertTrue(struct.getFields().get(0).getDataType().isVariantType()); + Assertions.assertTrue(namedStruct.getFields().get(0).getDataType().isVariantType()); + return true; + }) + ); + + AnalysisException mapKeyException = Assertions.assertThrowsExactly( + AnalysisException.class, + () -> PlanChecker.from(connectContext) + .analyze("select map(parse_to_variant('1'), 1)")); + Assertions.assertTrue(mapKeyException.getMessage() + .contains("map does not support variant keys")); + } + @Test public void testOverrideArity() { // the substring function has 2 override functions: diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/expression/check/CheckCastTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/expression/check/CheckCastTest.java index 2179b2d9ecb385..2c308b7baaa41e 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/expression/check/CheckCastTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/expression/check/CheckCastTest.java @@ -23,6 +23,7 @@ import org.apache.doris.nereids.types.BitmapType; import org.apache.doris.nereids.types.BooleanType; import org.apache.doris.nereids.types.CharType; +import org.apache.doris.nereids.types.ConnectorComputeVariantType; import org.apache.doris.nereids.types.DataType; import org.apache.doris.nereids.types.DateTimeType; import org.apache.doris.nereids.types.DateTimeV2Type; @@ -122,12 +123,26 @@ public void testTimeStampNsCastMatrix() { @Test public void testCastBetweenVariantTypes() { - VariantType v1Source = new VariantType(100); - VariantType v1SameProperties = new VariantType(100); - VariantType v1DifferentProperties = new VariantType(200); - - Assertions.assertTrue(CheckCast.check(v1Source, v1SameProperties, true)); - Assertions.assertTrue(CheckCast.check(v1Source, v1DifferentProperties, true)); + VariantType source = new VariantType(100); + VariantType sameProperties = new VariantType(100); + VariantType differentProperties = new VariantType(200); + + Assertions.assertTrue(CheckCast.check(source, sameProperties, true)); + Assertions.assertTrue(CheckCast.check(source, differentProperties, true)); + Assertions.assertTrue(CheckCast.check(source, ConnectorComputeVariantType.INSTANCE, true)); + Assertions.assertTrue(CheckCast.check(ConnectorComputeVariantType.INSTANCE, source, true)); + Assertions.assertTrue(CheckCast.check( + ConnectorComputeVariantType.INSTANCE, differentProperties, true)); + Assertions.assertTrue(CheckCast.check( + ArrayType.of(source), ArrayType.of(ConnectorComputeVariantType.INSTANCE), true)); + Assertions.assertTrue(CheckCast.check( + MapType.of(StringType.INSTANCE, source), + MapType.of(StringType.INSTANCE, ConnectorComputeVariantType.INSTANCE), true)); + Assertions.assertTrue(CheckCast.check( + new StructType(Lists.newArrayList(new StructField("v", source, true, ""))), + new StructType(Lists.newArrayList( + new StructField("v", ConnectorComputeVariantType.INSTANCE, true, ""))), + true)); } @Test diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/expression/rules/CaseWhenToCompoundPredicateTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/expression/rules/CaseWhenToCompoundPredicateTest.java index d15048122829d2..ea72fedd11f265 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/expression/rules/CaseWhenToCompoundPredicateTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/expression/rules/CaseWhenToCompoundPredicateTest.java @@ -20,12 +20,15 @@ import org.apache.doris.nereids.rules.expression.ExpressionRewriteContext; import org.apache.doris.nereids.rules.expression.ExpressionRewriteTestHelper; import org.apache.doris.nereids.rules.expression.ExpressionRuleExecutor; +import org.apache.doris.nereids.trees.expressions.functions.scalar.ShortCircuitIf; +import org.apache.doris.nereids.trees.expressions.literal.BooleanLiteral; import org.apache.doris.nereids.trees.plans.RelationId; import org.apache.doris.nereids.trees.plans.logical.LogicalEmptyRelation; import org.apache.doris.nereids.trees.plans.logical.LogicalFilter; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; +import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; class CaseWhenToCompoundPredicateTest extends ExpressionRewriteTestHelper { @@ -84,4 +87,16 @@ void testIfInCond() { context = oldContext; } } + + @Test + void testShortCircuitIfIsNotRewritten() { + executor = new ExpressionRuleExecutor(ImmutableList.of( + bottomUp(CaseWhenToCompoundPredicate.INSTANCE))); + ShortCircuitIf guarded = new ShortCircuitIf( + BooleanLiteral.TRUE, BooleanLiteral.TRUE, BooleanLiteral.FALSE); + Assertions.assertInstanceOf(ShortCircuitIf.class, executor.rewrite(guarded, context)); + + setExpressionOnFilter(); + Assertions.assertInstanceOf(ShortCircuitIf.class, executor.rewrite(guarded, context)); + } } diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ChangelogRowLevelDmlTransformTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ChangelogRowLevelDmlTransformTest.java new file mode 100644 index 00000000000000..687167136cd701 --- /dev/null +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ChangelogRowLevelDmlTransformTest.java @@ -0,0 +1,144 @@ +// 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. + +package org.apache.doris.nereids.trees.plans.commands; + +import org.apache.doris.analysis.UserIdentity; +import org.apache.doris.authorization.DataMaskSpec; +import org.apache.doris.catalog.Column; +import org.apache.doris.catalog.Env; +import org.apache.doris.catalog.ScalarType; +import org.apache.doris.datasource.CatalogIf; +import org.apache.doris.datasource.ExternalDatabase; +import org.apache.doris.datasource.plugin.PluginDrivenExternalTable; +import org.apache.doris.mysql.privilege.AccessControllerManager; +import org.apache.doris.nereids.analyzer.UnboundConnectorTableSink; +import org.apache.doris.nereids.exceptions.AnalysisException; +import org.apache.doris.nereids.trees.expressions.EqualTo; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.plans.JoinType; +import org.apache.doris.nereids.trees.plans.RelationId; +import org.apache.doris.nereids.trees.plans.commands.merge.MergeNotMatchedClause; +import org.apache.doris.nereids.trees.plans.logical.LogicalEmptyRelation; +import org.apache.doris.nereids.trees.plans.logical.LogicalJoin; +import org.apache.doris.nereids.trees.plans.logical.LogicalPlan; +import org.apache.doris.nereids.types.IntegerType; +import org.apache.doris.qe.ConnectContext; + +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentMatchers; +import org.mockito.Mockito; + +import java.util.Collections; +import java.util.Optional; + +public class ChangelogRowLevelDmlTransformTest { + + @Test + public void rejectsMaskedTargetForNonAdminUser() { + ConnectContext context = Mockito.mock(ConnectContext.class); + UserIdentity user = Mockito.mock(UserIdentity.class); + Env env = Mockito.mock(Env.class); + AccessControllerManager accessManager = Mockito.mock(AccessControllerManager.class); + PluginDrivenExternalTable table = Mockito.mock(PluginDrivenExternalTable.class); + ExternalDatabase database = Mockito.mock(ExternalDatabase.class); + CatalogIf catalog = Mockito.mock(CatalogIf.class); + + Mockito.when(context.getCurrentUserIdentity()).thenReturn(user); + Mockito.when(context.getEnv()).thenReturn(env); + Mockito.when(env.getAccessManager()).thenReturn(accessManager); + Mockito.when(table.getDatabase()).thenReturn(database); + Mockito.when(database.getCatalog()).thenReturn(catalog); + Mockito.when(catalog.getName()).thenReturn("catalog"); + Mockito.when(database.getFullName()).thenReturn("database"); + Mockito.when(table.getName()).thenReturn("target"); + Mockito.when(table.getFullSchema()).thenReturn( + ImmutableList.of(new Column("id", ScalarType.INT))); + Mockito.when(accessManager.evalDataMaskPolicies( + Mockito.eq(user), Mockito.eq("catalog"), Mockito.eq("database"), Mockito.eq("target"), + ArgumentMatchers.anySet())).thenReturn( + ImmutableMap.of("id", new DataMaskSpec("masked", "mask(id)"))); + + AnalysisException exception = Assertions.assertThrows(AnalysisException.class, + () -> ChangelogRowLevelDmlTransform.requireNoDataMask( + context, table, RowLevelDmlOp.UPDATE)); + Assertions.assertTrue(exception.getMessage().contains("data masking policies")); + } + + @Test + public void allowsUnmaskedTargetAndTrustedUsers() { + ConnectContext context = Mockito.mock(ConnectContext.class); + UserIdentity user = Mockito.mock(UserIdentity.class); + Env env = Mockito.mock(Env.class); + AccessControllerManager accessManager = Mockito.mock(AccessControllerManager.class); + PluginDrivenExternalTable table = Mockito.mock(PluginDrivenExternalTable.class); + ExternalDatabase database = Mockito.mock(ExternalDatabase.class); + CatalogIf catalog = Mockito.mock(CatalogIf.class); + + Mockito.when(context.getCurrentUserIdentity()).thenReturn(user); + Mockito.when(context.getEnv()).thenReturn(env); + Mockito.when(env.getAccessManager()).thenReturn(accessManager); + Mockito.when(table.getDatabase()).thenReturn(database); + Mockito.when(database.getCatalog()).thenReturn(catalog); + Mockito.when(catalog.getName()).thenReturn("catalog"); + Mockito.when(database.getFullName()).thenReturn("database"); + Mockito.when(table.getName()).thenReturn("target"); + Mockito.when(table.getFullSchema()).thenReturn( + ImmutableList.of(new Column("id", ScalarType.INT))); + Mockito.when(accessManager.evalDataMaskPolicies( + Mockito.eq(user), Mockito.eq("catalog"), Mockito.eq("database"), Mockito.eq("target"), + ArgumentMatchers.anySet())).thenReturn(Collections.emptyMap()); + + Assertions.assertDoesNotThrow(() -> ChangelogRowLevelDmlTransform.requireNoDataMask( + context, table, RowLevelDmlOp.DELETE)); + + Mockito.when(user.isRootUser()).thenReturn(true); + Assertions.assertDoesNotThrow(() -> ChangelogRowLevelDmlTransform.requireNoDataMask( + context, table, RowLevelDmlOp.MERGE)); + } + + @Test + public void mergeKeepsTargetOnProbeSide() { + ConnectContext context = Mockito.mock(ConnectContext.class); + UserIdentity user = Mockito.mock(UserIdentity.class); + PluginDrivenExternalTable table = Mockito.mock(PluginDrivenExternalTable.class); + Mockito.when(context.getCurrentUserIdentity()).thenReturn(user); + Mockito.when(user.isRootUser()).thenReturn(true); + + SlotReference sourceId = new SlotReference("id", IntegerType.INSTANCE, false, + ImmutableList.of("source")); + LogicalPlan source = new LogicalEmptyRelation(new RelationId(2), ImmutableList.of(sourceId)); + RowLevelDmlArgs args = RowLevelDmlArgs.forMerge(table, ImmutableList.of("catalog", "db", "target"), + Optional.empty(), Optional.empty(), source, + new EqualTo(new org.apache.doris.nereids.analyzer.UnboundSlot( + ImmutableList.of("target", "id")), sourceId), + ImmutableList.of(), ImmutableList.of(new MergeNotMatchedClause( + Optional.empty(), ImmutableList.of("id"), ImmutableList.of(sourceId)))); + + LogicalPlan result = new ChangelogRowLevelDmlTransform().synthesize(context, args, RowLevelDmlOp.MERGE); + + Assertions.assertInstanceOf(UnboundConnectorTableSink.class, result); + LogicalJoin join = (LogicalJoin) result.child(0); + Assertions.assertEquals(JoinType.RIGHT_OUTER_JOIN, join.getJoinType()); + Assertions.assertFalse(join.left() instanceof LogicalEmptyRelation, + "the target must stay on the probe/left side"); + Assertions.assertSame(source, join.right()); + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ExternalRowLevelMergePlanBuilderTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ExternalRowLevelMergePlanBuilderTest.java index 84509129231ccc..cae860ff25352d 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ExternalRowLevelMergePlanBuilderTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/ExternalRowLevelMergePlanBuilderTest.java @@ -29,6 +29,7 @@ import org.apache.doris.nereids.trees.expressions.Cast; import org.apache.doris.nereids.trees.expressions.EqualTo; import org.apache.doris.nereids.trees.expressions.NamedExpression; +import org.apache.doris.nereids.trees.expressions.functions.scalar.ShortCircuitIf; import org.apache.doris.nereids.trees.expressions.literal.BooleanLiteral; import org.apache.doris.nereids.trees.expressions.literal.IntegerLiteral; import org.apache.doris.nereids.trees.expressions.literal.StringLiteral; @@ -136,5 +137,7 @@ public void mergeBranchesUseThePinnedWriterType() { Assertions.assertEquals(uuidType, projections.get(0).child(0).getDataType(), "MERGE branch selection must be analyzable even when an expression starts with a wider type"); + Assertions.assertInstanceOf(ShortCircuitIf.class, projections.get(0).child(0), + "Iceberg MERGE must evaluate only the selected branch regardless of the session setting"); } } diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/IcebergRowLevelDmlTransformTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/PositionDeleteRowLevelDmlTransformTest.java similarity index 83% rename from fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/IcebergRowLevelDmlTransformTest.java rename to fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/PositionDeleteRowLevelDmlTransformTest.java index a079b576708d97..b838ad91ce61b7 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/IcebergRowLevelDmlTransformTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/PositionDeleteRowLevelDmlTransformTest.java @@ -29,6 +29,7 @@ import org.apache.doris.connector.spi.handle.ConnectorTransaction; import org.apache.doris.connector.spi.handle.WriteOperation; import org.apache.doris.connector.spi.pushdown.ConnectorPredicate; +import org.apache.doris.connector.spi.write.ConnectorRowChangeStyle; import org.apache.doris.datasource.ExternalDatabase; import org.apache.doris.datasource.plugin.PluginDrivenExternalCatalog; import org.apache.doris.datasource.plugin.PluginDrivenExternalTable; @@ -64,7 +65,7 @@ import java.util.Set; /** - * Unit tests for {@link IcebergRowLevelDmlTransform} (P6.3-T07c). + * Unit tests for {@link PositionDeleteRowLevelDmlTransform} (P6.3-T07c). * *

Covers the genuinely new T07c logic: the registry table-type predicate, the frozen per-op label * prefixes (profile/txn parity), the O5-2 synthetic-column exclusion supplied to {@link WriteConstraintExtractor @@ -72,10 +73,10 @@ * synthesis/executor/sink delegation had native end-to-end coverage in {@code IcebergDDLAndDMLPlanTest}, * retired with the P6.6 iceberg cutover (the native arm is no longer reachable).

*/ -public class IcebergRowLevelDmlTransformTest { +public class PositionDeleteRowLevelDmlTransformTest { private static final long TARGET_ID = 7L; - private final IcebergRowLevelDmlTransform transform = new IcebergRowLevelDmlTransform(); + private final PositionDeleteRowLevelDmlTransform transform = new PositionDeleteRowLevelDmlTransform(); private SlotReference slot(TableIf table, String name) { return SlotReference.fromColumn(StatementScopeIdGenerator.newExprId(), table, @@ -96,6 +97,11 @@ private Plan filterOver(TableIf table, String columnName) { * {@code getConnector().getWritePlanProvider(handle).supportedOperations()} probe. */ private static PluginDrivenExternalTable pluginTable(boolean supportsDelete, boolean supportsMerge) { + return pluginTable(supportsDelete, supportsMerge, ConnectorRowChangeStyle.POSITION_DELETE); + } + + private static PluginDrivenExternalTable pluginTable(boolean supportsDelete, boolean supportsMerge, + ConnectorRowChangeStyle style) { PluginDrivenExternalTable table = Mockito.mock(PluginDrivenExternalTable.class); PluginDrivenExternalCatalog catalog = Mockito.mock(PluginDrivenExternalCatalog.class); Connector connector = Mockito.mock(Connector.class); @@ -108,6 +114,7 @@ private static PluginDrivenExternalTable pluginTable(boolean supportsDelete, boo } Mockito.when(table.getCatalog()).thenReturn(catalog); Mockito.when(catalog.getConnector()).thenReturn(connector); + Mockito.when(table.getConnectorRowChangeStyle()).thenReturn(style); // The row-level DML admission probe now resolves per-handle via the table helper; stub it directly. The // catalog -> connector chain is still needed for checkMode (validateRowLevelDmlMode). Mockito.when(table.connectorSupportedWriteOperations()).thenReturn(ops); @@ -116,22 +123,40 @@ private static PluginDrivenExternalTable pluginTable(boolean supportsDelete, boo @Test public void handlesPluginDrivenTableByRowLevelDmlCapability() { - // An iceberg table presents as PluginDrivenExternalTable; it is admitted via the - // neutral connector capability (supportsDelete || supportsMerge), NOT a concrete iceberg cast. + // Position-delete tables are admitted by representation and capability, not a concrete source cast. Assertions.assertTrue(transform.handles(pluginTable(true, false))); Assertions.assertTrue(transform.handles(pluginTable(false, true))); Assertions.assertTrue(transform.handles(pluginTable(true, true))); - // A plugin connector with neither capability (e.g. jdbc/es/paimon today) must NOT be admitted, + PluginDrivenExternalTable updateOnly = pluginTable(false, false); + Mockito.when(updateOnly.connectorSupportedWriteOperations()) + .thenReturn(EnumSet.of(WriteOperation.UPDATE)); + Assertions.assertTrue(transform.handles(updateOnly)); + // A plugin connector with neither capability (e.g. jdbc/es) must NOT be admitted, // else its row-level DML would route through the iceberg synthesis path. Assertions.assertFalse(transform.handles(pluginTable(false, false))); // Non-plugin table types and null are never admitted. Assertions.assertFalse(transform.handles(Mockito.mock(TableIf.class))); Assertions.assertFalse(transform.handles(null)); + Assertions.assertTrue(transform.requiresExternalTableBatchModeDisabled()); + } + + @Test + public void registryRoutesEachRowChangeRepresentationToItsTransform() { + PluginDrivenExternalTable changelog = pluginTable(true, true, ConnectorRowChangeStyle.CHANGELOG); + PluginDrivenExternalTable undeclared = pluginTable(true, false, ConnectorRowChangeStyle.NONE); + + Assertions.assertFalse(transform.handles(changelog)); + Assertions.assertFalse(transform.handles(undeclared)); + Assertions.assertTrue(RowLevelDmlRegistry.find(changelog) + .orElseThrow(AssertionError::new) instanceof ChangelogRowLevelDmlTransform); + Assertions.assertThrows(AnalysisException.class, () -> RowLevelDmlRegistry.find(undeclared)); + Assertions.assertTrue(RowLevelDmlRegistry.find(pluginTable(true, true)) + .orElseThrow(AssertionError::new) instanceof PositionDeleteRowLevelDmlTransform); } /** * A {@link PluginDrivenExternalTable} (db1.t1) whose connector resolves to {@code metadata}. Used to - * drive the post-flip {@link IcebergRowLevelDmlTransform#checkMode} plugin arm, which routes the + * drive the post-flip {@link PositionDeleteRowLevelDmlTransform#checkMode} plugin arm, which routes the * copy-on-write rejection through the connector's neutral {@code validateRowLevelDmlMode} SPI. */ private static PluginDrivenExternalTable pluginTableWithMetadata( @@ -146,6 +171,8 @@ private static PluginDrivenExternalTable pluginTableWithMetadata( Mockito.when(catalog.buildConnectorSession()).thenReturn(session); Mockito.when(catalog.getConnector()).thenReturn(connector); Mockito.when(connector.getMetadata(session)).thenReturn(metadata); + Mockito.when(table.connectorSupportedWriteOperations()).thenReturn( + EnumSet.of(WriteOperation.DELETE, WriteOperation.UPDATE, WriteOperation.MERGE)); // checkMode now resolves metadata through the per-statement funnel, which reads the session's statement // scope; offline tests use NONE (a fresh getMetadata per call, byte-identical to pre-funnel). Mockito.when(session.getStatementScope()).thenReturn(ConnectorStatementScope.NONE); @@ -237,20 +264,6 @@ public void synthesizeDeleteOnPluginTableBuildsSinkTargetingIt() { Mockito.verify(table, Mockito.times(1)).getWriteSchemaSnapshot(); } - @Test - public void setupConflictDetectionPluginArmIsNoOp() { - // The conflict filter runs ONLY through the SPI path (applyWriteConstraintIfPresent), so - // setupConflictDetection is a no-op that must NOT touch the executor (the retired native arm cast - // it to Iceberg{Delete,Merge}Executor and called setConflictDetectionFilter). - BaseExternalTableInsertExecutor executor = Mockito.mock(PluginDrivenInsertExecutor.class); - PluginDrivenExternalTable table = Mockito.mock(PluginDrivenExternalTable.class); - Plan analyzedPlan = Mockito.mock(Plan.class); - - Assertions.assertDoesNotThrow(() -> - transform.setupConflictDetection(executor, analyzedPlan, table, RowLevelDmlOp.DELETE)); - Mockito.verifyNoInteractions(executor); - } - @Test public void finalizeSinkPluginArmRoutesToConnectorFinalize() { // Finalize goes through the connector's single transaction model (no rewritable-delete @@ -268,15 +281,24 @@ public void finalizeSinkPluginArmRoutesToConnectorFinalize() { @Test public void labelPrefixIsFrozenPerOp() { // These are profile/txn-visible and must stay byte-identical to the legacy Iceberg*Command labels. - Assertions.assertEquals("iceberg_delete", transform.labelPrefix(RowLevelDmlOp.DELETE)); - Assertions.assertEquals("iceberg_update_merge", transform.labelPrefix(RowLevelDmlOp.UPDATE)); - Assertions.assertEquals("iceberg_merge_into", transform.labelPrefix(RowLevelDmlOp.MERGE)); + PluginDrivenExternalTable table = Mockito.mock(PluginDrivenExternalTable.class); + Mockito.when(table.getConnectorRowLevelDmlLabelPrefix(WriteOperation.DELETE)) + .thenReturn("iceberg_delete"); + Mockito.when(table.getConnectorRowLevelDmlLabelPrefix(WriteOperation.UPDATE)) + .thenReturn("iceberg_update_merge"); + Mockito.when(table.getConnectorRowLevelDmlLabelPrefix(WriteOperation.MERGE)) + .thenReturn("iceberg_merge_into"); + Assertions.assertEquals("iceberg_delete", transform.labelPrefix(table, RowLevelDmlOp.DELETE)); + Assertions.assertEquals("iceberg_update_merge", transform.labelPrefix(table, RowLevelDmlOp.UPDATE)); + Assertions.assertEquals("iceberg_merge_into", transform.labelPrefix(table, RowLevelDmlOp.MERGE)); } @Test public void extractWriteConstraintKeepsRegularTargetColumn() { - TableIf target = Mockito.mock(PluginDrivenExternalTable.class); + PluginDrivenExternalTable target = Mockito.mock(PluginDrivenExternalTable.class); Mockito.when(target.getId()).thenReturn(TARGET_ID); + Mockito.when(target.getConnectorRowLevelWriteConstraintExcludedColumns()) + .thenReturn(ImmutableSet.of(Column.ICEBERG_ROWID_COL)); Optional result = transform.extractWriteConstraint(filterOver(target, "id"), target); Assertions.assertTrue(result.isPresent()); } @@ -285,16 +307,20 @@ public void extractWriteConstraintKeepsRegularTargetColumn() { public void extractWriteConstraintExcludesRowIdColumn() { // Load-bearing: the synthetic $row_id slot has originalTable == target, so the origin-table check alone // would keep it; only the iceberg ICEBERG_EXCLUSION predicate drops it (closes T07b critic BLOCKER). - TableIf target = Mockito.mock(PluginDrivenExternalTable.class); + PluginDrivenExternalTable target = Mockito.mock(PluginDrivenExternalTable.class); Mockito.when(target.getId()).thenReturn(TARGET_ID); + Mockito.when(target.getConnectorRowLevelWriteConstraintExcludedColumns()) + .thenReturn(ImmutableSet.of(Column.ICEBERG_ROWID_COL)); Plan plan = filterOver(target, Column.ICEBERG_ROWID_COL); Assertions.assertFalse(transform.extractWriteConstraint(plan, target).isPresent()); } @Test public void extractWriteConstraintExcludesMetadataColumn() { - TableIf target = Mockito.mock(PluginDrivenExternalTable.class); + PluginDrivenExternalTable target = Mockito.mock(PluginDrivenExternalTable.class); Mockito.when(target.getId()).thenReturn(TARGET_ID); + Mockito.when(target.getConnectorRowLevelWriteConstraintExcludedColumns()) + .thenReturn(ImmutableSet.of("$partition_spec_id")); // "$partition_spec_id" is a position-delete metadata column -> excluded. Plan plan = filterOver(target, "$partition_spec_id"); Assertions.assertFalse(transform.extractWriteConstraint(plan, target).isPresent()); diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/RowLevelDmlRowIdUtilsTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/RowLevelDmlRowIdUtilsTest.java index 04542cbcf12760..a0d6895c89d5cb 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/RowLevelDmlRowIdUtilsTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/RowLevelDmlRowIdUtilsTest.java @@ -78,6 +78,14 @@ public void isRowIdInjectionTargetAcceptsMergeOnlyPluginDrivenTable() { RowLevelDmlRowIdUtils.isRowIdInjectionTarget(pluginTableWithCapability(false, true))); } + @Test + public void isRowIdInjectionTargetAcceptsUpdateOnlyPluginDrivenTable() { + PluginDrivenExternalTable table = pluginTableWithCapability(false, false); + Mockito.when(table.connectorSupportedWriteOperations()).thenReturn(EnumSet.of(WriteOperation.UPDATE)); + + Assertions.assertTrue(RowLevelDmlRowIdUtils.isRowIdInjectionTarget(table)); + } + @Test public void isRowIdInjectionTargetRejectsPluginDrivenTableWithoutCapability() { // A non-iceberg plugin-driven table (jdbc/es/trino/max_compute/paimon) declares neither capability, diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/merge/MergeIntoCommandTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/merge/MergeIntoCommandTest.java index b2d9b32b708bcb..0d7bed77b2b7e3 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/merge/MergeIntoCommandTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/commands/merge/MergeIntoCommandTest.java @@ -37,7 +37,6 @@ import org.apache.doris.nereids.trees.expressions.IsNull; import org.apache.doris.nereids.trees.expressions.NamedExpression; import org.apache.doris.nereids.trees.expressions.Not; -import org.apache.doris.nereids.trees.expressions.functions.scalar.If; import org.apache.doris.nereids.trees.expressions.literal.IntegerLiteral; import org.apache.doris.nereids.trees.expressions.literal.NullLiteral; import org.apache.doris.nereids.trees.expressions.literal.TinyIntLiteral; @@ -167,11 +166,11 @@ public void testGenerateBranchLabel() throws Exception { Method generateBranchLabel = clazz.getDeclaredMethod("generateBranchLabel", NamedExpression.class); generateBranchLabel.setAccessible(true); NamedExpression result = (NamedExpression) generateBranchLabel.invoke(command, unboundSlot); - Expression matchedLabel = new If(new IntegerLiteral(1), new IntegerLiteral(0), - new If(new IntegerLiteral(2), new IntegerLiteral(1), new IntegerLiteral(2))); - Expression notMatchedLabel = new If(new IntegerLiteral(3), new IntegerLiteral(3), - new If(new IntegerLiteral(4), new IntegerLiteral(4), new IntegerLiteral(5))); - NamedExpression expected = new UnboundAlias(new If(new Not(new IsNull(unboundSlot)), + Expression matchedLabel = MergeUtils.selectBranch(new IntegerLiteral(1), new IntegerLiteral(0), + MergeUtils.selectBranch(new IntegerLiteral(2), new IntegerLiteral(1), new IntegerLiteral(2))); + Expression notMatchedLabel = MergeUtils.selectBranch(new IntegerLiteral(3), new IntegerLiteral(3), + MergeUtils.selectBranch(new IntegerLiteral(4), new IntegerLiteral(4), new IntegerLiteral(5))); + NamedExpression expected = new UnboundAlias(MergeUtils.selectBranch(new Not(new IsNull(unboundSlot)), matchedLabel, notMatchedLabel), "__DORIS_MERGE_INTO_BRANCH_LABEL__"); Assertions.assertEquals(expected, result); } @@ -943,14 +942,30 @@ public void testGenerateFinalProjections() throws Exception { generateFinalProjections.setAccessible(true); List result = (List) generateFinalProjections.invoke(command, colNames, finalProjections); List expected = ImmutableList.of( - new UnboundAlias(new If(new EqualTo(new UnboundSlot("__DORIS_MERGE_INTO_BRANCH_LABEL__"), new IntegerLiteral(3)), new IntegerLiteral(41), - new If(new EqualTo(new UnboundSlot("__DORIS_MERGE_INTO_BRANCH_LABEL__"), new IntegerLiteral(2)), new IntegerLiteral(31), - new If(new EqualTo(new UnboundSlot("__DORIS_MERGE_INTO_BRANCH_LABEL__"), new IntegerLiteral(1)), new IntegerLiteral(21), - new If(new EqualTo(new UnboundSlot("__DORIS_MERGE_INTO_BRANCH_LABEL__"), new IntegerLiteral(0)), new IntegerLiteral(11), new NullLiteral())))), "c1"), - new UnboundAlias(new If(new EqualTo(new UnboundSlot("__DORIS_MERGE_INTO_BRANCH_LABEL__"), new IntegerLiteral(3)), new IntegerLiteral(42), - new If(new EqualTo(new UnboundSlot("__DORIS_MERGE_INTO_BRANCH_LABEL__"), new IntegerLiteral(2)), new IntegerLiteral(32), - new If(new EqualTo(new UnboundSlot("__DORIS_MERGE_INTO_BRANCH_LABEL__"), new IntegerLiteral(1)), new IntegerLiteral(22), - new If(new EqualTo(new UnboundSlot("__DORIS_MERGE_INTO_BRANCH_LABEL__"), new IntegerLiteral(0)), new IntegerLiteral(12), new NullLiteral())))), "c2") + new UnboundAlias(MergeUtils.selectBranch( + new EqualTo(new UnboundSlot("__DORIS_MERGE_INTO_BRANCH_LABEL__"), new IntegerLiteral(3)), + new IntegerLiteral(41), MergeUtils.selectBranch( + new EqualTo(new UnboundSlot("__DORIS_MERGE_INTO_BRANCH_LABEL__"), + new IntegerLiteral(2)), + new IntegerLiteral(31), MergeUtils.selectBranch( + new EqualTo(new UnboundSlot("__DORIS_MERGE_INTO_BRANCH_LABEL__"), + new IntegerLiteral(1)), + new IntegerLiteral(21), MergeUtils.selectBranch( + new EqualTo(new UnboundSlot("__DORIS_MERGE_INTO_BRANCH_LABEL__"), + new IntegerLiteral(0)), + new IntegerLiteral(11), new NullLiteral())))), "c1"), + new UnboundAlias(MergeUtils.selectBranch( + new EqualTo(new UnboundSlot("__DORIS_MERGE_INTO_BRANCH_LABEL__"), new IntegerLiteral(3)), + new IntegerLiteral(42), MergeUtils.selectBranch( + new EqualTo(new UnboundSlot("__DORIS_MERGE_INTO_BRANCH_LABEL__"), + new IntegerLiteral(2)), + new IntegerLiteral(32), MergeUtils.selectBranch( + new EqualTo(new UnboundSlot("__DORIS_MERGE_INTO_BRANCH_LABEL__"), + new IntegerLiteral(1)), + new IntegerLiteral(22), MergeUtils.selectBranch( + new EqualTo(new UnboundSlot("__DORIS_MERGE_INTO_BRANCH_LABEL__"), + new IntegerLiteral(0)), + new IntegerLiteral(12), new NullLiteral())))), "c2") ); Assertions.assertEquals(expected, result); diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/physical/PhysicalConnectorTableSinkTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/physical/PhysicalConnectorTableSinkTest.java index 22a7b4aca04bb3..eebe093cedeeab 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/physical/PhysicalConnectorTableSinkTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/physical/PhysicalConnectorTableSinkTest.java @@ -20,7 +20,10 @@ import org.apache.doris.catalog.Column; import org.apache.doris.catalog.PrimitiveType; import org.apache.doris.common.jmockit.Deencapsulation; +import org.apache.doris.connector.spi.write.ConnectorWriteDistribution; import org.apache.doris.datasource.plugin.PluginDrivenExternalTable; +import org.apache.doris.nereids.properties.DistributionSpecExternalTableSinkHashPartitioned; +import org.apache.doris.nereids.properties.DistributionSpecHash; import org.apache.doris.nereids.properties.DistributionSpecHiveTableSinkHashPartitioned; import org.apache.doris.nereids.properties.MustLocalSortOrderSpec; import org.apache.doris.nereids.properties.OrderKey; @@ -37,6 +40,7 @@ import java.util.Arrays; import java.util.List; +import java.util.Optional; /** * Tests for {@link PhysicalConnectorTableSink#getRequirePhysicalProperties()} (FIX-WRITE-DISTRIBUTION, @@ -305,6 +309,24 @@ public void partitionHashWriteHashesByPartitionWithoutLocalSort() { + "would pay an unnecessary sort the legacy path never had"); } + @Test + public void changelogWriteSkipsOperationColumnWhenLocatingPartition() { + SlotReference operationSlot = new SlotReference("connector_operation", IntegerType.INSTANCE); + SlotReference dataSlot = new SlotReference("data", IntegerType.INSTANCE); + SlotReference partSlot = new SlotReference("part", IntegerType.INSTANCE); + PhysicalConnectorTableSink sink = sink( + table(true, false, true, ImmutableList.of(PART), ImmutableList.of(DATA, PART)), + Arrays.asList(DATA, PART), + ImmutableList.of(operationSlot, dataSlot, partSlot)); + Deencapsulation.setField(sink, "hasRowOperationColumn", true); + + PhysicalProperties props = sink.getRequirePhysicalProperties(); + DistributionSpecHiveTableSinkHashPartitioned dist = + (DistributionSpecHiveTableSinkHashPartitioned) props.getDistributionSpec(); + Assertions.assertEquals(ImmutableList.of(partSlot.getExprId()), dist.getOutputColExprIds(), + "the connector operation column must not shift partition routing onto a data column"); + } + /** * Non-partitioned write on a hash-write connector: the hash arm's {@code !partitionNames.isEmpty()} * gate falls through to the parallel arm, matching legacy {@code PhysicalHiveTableSink}'s @@ -323,6 +345,68 @@ public void nonPartitionedHashWriteConnectorUsesRandomWhenParallel() { "a non-partitioned hash-write connector falls through to parallel writers, not the hash arm"); } + @Test + public void connectorOwnedHashDistributionStaysOpaqueInFeCore() { + SlotReference operationSlot = new SlotReference("connector_operation", IntegerType.INSTANCE); + SlotReference dataSlot = new SlotReference("data", IntegerType.INSTANCE); + SlotReference partSlot = new SlotReference("part", IntegerType.INSTANCE); + PluginDrivenExternalTable table = table(false, false, ImmutableList.of(), + ImmutableList.of(DATA, PART)); + Mockito.when(table.getConnectorWriteDistribution()).thenReturn(Optional.of( + ConnectorWriteDistribution.externalHash(ImmutableList.of("part"), + "connector_bucket", java.util.Collections.singletonMap("buckets", "8"), + ConnectorWriteDistribution.WriterAssignment.IDENTITY))); + PhysicalConnectorTableSink sink = sink(table, Arrays.asList(DATA, PART), + ImmutableList.of(operationSlot, dataSlot, partSlot)); + Deencapsulation.setField(sink, "hasRowOperationColumn", true); + + DistributionSpecExternalTableSinkHashPartitioned distribution + = (DistributionSpecExternalTableSinkHashPartitioned) + sink.getRequirePhysicalProperties().getDistributionSpec(); + Assertions.assertEquals(ImmutableList.of(partSlot.getExprId()), + distribution.getOutputColumnExprIds()); + Assertions.assertEquals("connector_bucket", distribution.getPartitionFunction()); + Assertions.assertEquals("8", distribution.getPartitionFunctionOptions().get("buckets")); + } + + @Test + public void connectorDistributionKeepsRequiredPartitionLocalSort() { + SlotReference dataSlot = new SlotReference("data", IntegerType.INSTANCE); + SlotReference partSlot = new SlotReference("part", IntegerType.INSTANCE); + PluginDrivenExternalTable table = table(true, true, ImmutableList.of(PART), + ImmutableList.of(DATA, PART)); + Mockito.when(table.getConnectorWriteDistribution()).thenReturn(Optional.of( + ConnectorWriteDistribution.externalHash(ImmutableList.of("part"), + "connector_bucket", java.util.Collections.emptyMap(), + ConnectorWriteDistribution.WriterAssignment.IDENTITY))); + PhysicalConnectorTableSink sink = sink(table, Arrays.asList(DATA, PART), + ImmutableList.of(dataSlot, partSlot)); + + PhysicalProperties properties = sink.getRequirePhysicalProperties(); + + Assertions.assertInstanceOf(DistributionSpecExternalTableSinkHashPartitioned.class, + properties.getDistributionSpec()); + Assertions.assertInstanceOf(MustLocalSortOrderSpec.class, properties.getOrderSpec()); + Assertions.assertEquals(partSlot, properties.getOrderSpec().getOrderKeys().get(0).getExpr()); + } + + @Test + public void nameMappedDistributionUsesExplicitColumnOrder() { + SlotReference partSlot = new SlotReference("part", IntegerType.INSTANCE); + SlotReference dataSlot = new SlotReference("data", IntegerType.INSTANCE); + PluginDrivenExternalTable table = table(false, false, ImmutableList.of(), + ImmutableList.of(DATA, PART)); + Mockito.when(table.getConnectorWriteDistribution()).thenReturn(Optional.of( + ConnectorWriteDistribution.hash(ImmutableList.of("part")))); + PhysicalConnectorTableSink sink = sink(table, Arrays.asList(PART, DATA), + ImmutableList.of(partSlot, dataSlot)); + + PhysicalProperties properties = sink.getRequirePhysicalProperties(); + + Assertions.assertEquals(ImmutableList.of(partSlot.getExprId()), + ((DistributionSpecHash) properties.getDistributionSpec()).getOrderedShuffledColumns()); + } + // ==================== helpers ==================== private static PluginDrivenExternalTable table(boolean parallelWrite, boolean requirePartitionSort, @@ -332,6 +416,7 @@ private static PluginDrivenExternalTable table(boolean parallelWrite, boolean re Mockito.when(table.requirePartitionLocalSortOnWrite()).thenReturn(requirePartitionSort); Mockito.when(table.getPartitionColumns()).thenReturn(partitionColumns); Mockito.when(table.getFullSchema()).thenReturn(fullSchema); + Mockito.when(table.getConnectorWriteDistribution()).thenReturn(Optional.empty()); return table; } diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/AbstractJobProcessorTest.java b/fe/fe-core/src/test/java/org/apache/doris/qe/AbstractJobProcessorTest.java index 64f57cfec1e111..d20ad967a4ed96 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/qe/AbstractJobProcessorTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/qe/AbstractJobProcessorTest.java @@ -19,15 +19,21 @@ import org.apache.doris.common.Status; import org.apache.doris.nereids.trees.plans.distribute.worker.BackendWorker; +import org.apache.doris.qe.runtime.BackendFragmentId; import org.apache.doris.qe.runtime.MultiFragmentsPipelineTask; import org.apache.doris.qe.runtime.PipelineExecutionTask; import org.apache.doris.qe.runtime.SingleFragmentPipelineTask; import org.apache.doris.thrift.TReportExecStatusParams; +import org.apache.doris.thrift.TStatus; +import org.apache.doris.thrift.TStatusCode; +import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; import org.mockito.Mockito; +import java.nio.ByteBuffer; import java.util.Collections; +import java.util.Map; import java.util.Optional; class AbstractJobProcessorTest { @@ -59,6 +65,18 @@ void fragmentDispatchBeforeFinishBroadcastsWhenExecutionFinishes() { Mockito.verify(fragmentsTask).cancelExecute(Status.FINISHED); } + @Test + void opaqueConnectorDataRequiresARegisteredFragmentHandler() { + TestJobProcessor processor = new TestJobProcessor(Mockito.mock(CoordinatorContext.class)); + processor.setBackendFragmentTasks(Collections.emptyMap()); + TReportExecStatusParams params = new TReportExecStatusParams() + .setStatus(new TStatus(TStatusCode.OK)) + .setConnectorCommitData(Collections.singletonList(ByteBuffer.wrap(new byte[] {1}))); + + Assertions.assertThrows(IllegalStateException.class, + () -> processor.updateFragmentExecStatus(params)); + } + private static TestJobProcessor createProcessor(MultiFragmentsPipelineTask fragmentsTask) { BackendWorker worker = Mockito.mock(BackendWorker.class); PipelineExecutionTask executionTask = Mockito.mock(PipelineExecutionTask.class); @@ -78,6 +96,10 @@ void setExecutionTask(PipelineExecutionTask executionTask) { this.executionTask = Optional.of(executionTask); } + void setBackendFragmentTasks(Map tasks) { + this.backendFragmentTasks = Optional.of(tasks); + } + @Override protected void doProcessReportExecStatus( TReportExecStatusParams params, SingleFragmentPipelineTask fragmentTask) {} diff --git a/fe/fe-core/src/test/java/org/apache/doris/qe/QeProcessorImplReportAckTest.java b/fe/fe-core/src/test/java/org/apache/doris/qe/QeProcessorImplReportAckTest.java index 2e15875c73ffd0..ad268080e20b57 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/qe/QeProcessorImplReportAckTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/qe/QeProcessorImplReportAckTest.java @@ -37,6 +37,7 @@ import java.lang.reflect.Field; import java.lang.reflect.Modifier; +import java.nio.ByteBuffer; import java.util.Collections; class QeProcessorImplReportAckTest { @@ -57,6 +58,18 @@ void rejectsExternalReportWithoutCoordinator() { Assertions.assertFalse(result.isExternalFileCommitDataAccepted()); } + @Test + void rejectsOpaqueConnectorReportWithoutCoordinator() { + TReportExecStatusParams params = params(new TUniqueId(12345, 6)); + params.unsetIcebergCommitDatas(); + params.setConnectorCommitData(Collections.singletonList(ByteBuffer.wrap(new byte[] {1}))); + + TReportExecStatusResult result = report(params); + + Assertions.assertEquals(TStatusCode.INTERNAL_ERROR, result.getStatus().getStatusCode()); + Assertions.assertFalse(result.isExternalFileCommitDataAccepted()); + } + @Test void rejectsExternalReportWhenHandlerThrows() throws Exception { TUniqueId queryId = new TUniqueId(12345, 2); @@ -99,6 +112,26 @@ void retriesAcceptedExternalReportAfterCoordinatorRemoval() throws Exception { Mockito.verify(coordinator, Mockito.times(1)).updateFragmentExecStatus(params); } + @Test + void retriesAcceptedOpaqueConnectorReportAfterCoordinatorRemoval() throws Exception { + TUniqueId queryId = new TUniqueId(12345, 7); + Coordinator coordinator = register(queryId); + Mockito.when(coordinator.updateFragmentExecStatus(Mockito.any())).thenReturn(true); + TReportExecStatusParams params = params(queryId); + params.unsetIcebergCommitDatas(); + params.setConnectorCommitData(Collections.singletonList(ByteBuffer.wrap(new byte[] {1}))); + + TReportExecStatusResult first = report(params); + QeProcessorImpl.INSTANCE.unregisterQuery(queryId); + registeredQueryId = null; + TReportExecStatusResult retry = report(params); + + Assertions.assertTrue(first.isExternalFileCommitDataAccepted()); + Assertions.assertTrue(retry.isExternalFileCommitDataAccepted()); + Assertions.assertEquals(TStatusCode.OK, retry.getStatus().getStatusCode()); + Mockito.verify(coordinator, Mockito.times(1)).updateFragmentExecStatus(params); + } + @Test void evictedAcceptanceTokenRejectsRetryAfterCoordinatorRemoval() throws Exception { TUniqueId queryId = new TUniqueId(12345, 5); diff --git a/fe/fe-core/src/test/java/org/apache/doris/transaction/CommitDataSerializerTest.java b/fe/fe-core/src/test/java/org/apache/doris/transaction/CommitDataSerializerTest.java index 3068a4a802c343..d30bdd5b44f599 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/transaction/CommitDataSerializerTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/transaction/CommitDataSerializerTest.java @@ -21,6 +21,7 @@ import org.apache.doris.thrift.THivePartitionUpdate; import org.apache.doris.thrift.TIcebergCommitData; import org.apache.doris.thrift.TMCCommitData; +import org.apache.doris.thrift.TReportExecStatusParams; import org.apache.doris.thrift.TUpdateMode; import org.apache.thrift.TBase; @@ -30,6 +31,7 @@ import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.Arrays; import java.util.List; @@ -137,4 +139,63 @@ public void addCommitData(byte[] commitFragment) { } } + @Test + public void rawFeedPreservesEachBinarySlice() { + List payloads = new ArrayList<>(); + Transaction collector = new Transaction() { + @Override + public void commit() { + throw new UnsupportedOperationException("commit not expected in this test"); + } + + @Override + public void rollback() { + throw new UnsupportedOperationException("rollback not expected in this test"); + } + + @Override + public void addCommitData(byte[] commitFragment) { + payloads.add(commitFragment); + } + }; + ByteBuffer fragment = ByteBuffer.wrap(new byte[] {0, 1, 2, 3}); + fragment.position(1); + fragment.limit(3); + + CommitDataSerializer.feedRaw(collector, Arrays.asList(fragment, ByteBuffer.wrap(new byte[] {4, 5}))); + + Assertions.assertArrayEquals(new byte[] {1, 2}, payloads.get(0)); + Assertions.assertArrayEquals(new byte[] {4, 5}, payloads.get(1)); + Assertions.assertEquals(1, fragment.position()); + } + + @Test + public void reportFeedRecognizesAndDeliversOpaqueConnectorData() { + List payloads = new ArrayList<>(); + Transaction collector = new Transaction() { + @Override + public void commit() { + throw new UnsupportedOperationException("commit not expected in this test"); + } + + @Override + public void rollback() { + throw new UnsupportedOperationException("rollback not expected in this test"); + } + + @Override + public void addCommitData(byte[] commitFragment) { + payloads.add(commitFragment); + } + }; + TReportExecStatusParams report = new TReportExecStatusParams() + .setConnectorCommitData(Arrays.asList(ByteBuffer.wrap(new byte[] {6, 7}))); + + Assertions.assertTrue(CommitDataSerializer.hasCommitData(report)); + CommitDataSerializer.feed(collector, report); + + Assertions.assertEquals(1, payloads.size()); + Assertions.assertArrayEquals(new byte[] {6, 7}, payloads.get(0)); + } + } diff --git a/fe/fe-core/src/test/java/org/apache/doris/transaction/PluginDrivenTransactionManagerTest.java b/fe/fe-core/src/test/java/org/apache/doris/transaction/PluginDrivenTransactionManagerTest.java index e9b2b5b6f39a18..dce2fc94e0f27e 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/transaction/PluginDrivenTransactionManagerTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/transaction/PluginDrivenTransactionManagerTest.java @@ -27,6 +27,9 @@ import java.util.ArrayList; import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; /** * Delegation tests for {@link PluginDrivenTransactionManager} and its internal @@ -107,6 +110,37 @@ public long allocateWriteBlockRange(String writeSessionId, long count) { } } + private static final class BlockingCommitDataTransaction extends RecordingConnectorTransaction { + private final CountDownLatch addEntered = new CountDownLatch(1); + private final CountDownLatch releaseAdd = new CountDownLatch(1); + private final CountDownLatch commitEntered = new CountDownLatch(1); + + private BlockingCommitDataTransaction(long txnId) { + super(txnId); + } + + @Override + public void addCommitData(byte[] commitFragment) { + addEntered.countDown(); + await(releaseAdd); + super.addCommitData(commitFragment); + } + + @Override + public void commit() { + commitEntered.countDown(); + } + + private static void await(CountDownLatch latch) { + try { + latch.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException(e); + } + } + } + @Test public void addCommitDataIsDelegatedToConnectorTransaction() throws UserException { PluginDrivenTransactionManager manager = new PluginDrivenTransactionManager(); @@ -116,7 +150,8 @@ public void addCommitDataIsDelegatedToConnectorTransaction() throws UserExceptio byte[] fragment = {1, 2, 3}; manager.getTransaction(txnId).addCommitData(fragment); - Assertions.assertEquals(1, connectorTx.commitFragments.size()); + Assertions.assertEquals(1, + ((RecordingConnectorTransaction) connectorTx).commitFragments.size()); Assertions.assertSame(fragment, connectorTx.commitFragments.get(0)); } @@ -182,6 +217,46 @@ public void legacyMarkerKeepsInertWriteDefaults() throws UserException { Assertions.assertFalse(txn instanceof WriteBlockAllocatingTransaction); } + @Test + public void commitWaitsForInFlightCommitDataAndRejectsLateReports() throws Exception { + PluginDrivenTransactionManager manager = new PluginDrivenTransactionManager(); + BlockingCommitDataTransaction connectorTx = new BlockingCommitDataTransaction(90000L); + long txnId = manager.begin(connectorTx); + Transaction transaction = manager.getTransaction(txnId); + AtomicReference failure = new AtomicReference<>(); + Thread reportThread = new Thread(() -> { + try { + transaction.addCommitData(new byte[] {1}); + } catch (Throwable t) { + failure.compareAndSet(null, t); + } + }); + Thread commitThread = new Thread(() -> { + try { + manager.commit(txnId); + } catch (Throwable t) { + failure.compareAndSet(null, t); + } + }); + + reportThread.start(); + Assertions.assertTrue(connectorTx.addEntered.await(5, TimeUnit.SECONDS)); + commitThread.start(); + Assertions.assertFalse(connectorTx.commitEntered.await(100, TimeUnit.MILLISECONDS), + "commit must not overlap addCommitData on a connector transaction"); + connectorTx.releaseAdd.countDown(); + reportThread.join(5000); + commitThread.join(5000); + + Assertions.assertFalse(reportThread.isAlive()); + Assertions.assertFalse(commitThread.isAlive()); + Assertions.assertNull(failure.get()); + Assertions.assertEquals(1, + ((RecordingConnectorTransaction) connectorTx).commitFragments.size()); + Assertions.assertThrows(IllegalStateException.class, + () -> transaction.addCommitData(new byte[] {2})); + } + // ──────────── global registration (P4-T06a W-d / gap G3) ──────────── // // begin(ConnectorTransaction) must also register the txn in the process-wide diff --git a/gensrc/thrift/FrontendService.thrift b/gensrc/thrift/FrontendService.thrift index fb3f4c3860b8ba..25166c50b70a54 100644 --- a/gensrc/thrift/FrontendService.thrift +++ b/gensrc/thrift/FrontendService.thrift @@ -338,6 +338,9 @@ struct TReportExecStatusParams { 32: optional list mc_commit_datas 33: optional string first_error_msg + + // Opaque, connector-owned commit fragments; FE routes them to the transaction. + 34: optional list connector_commit_data } struct TFeResult { diff --git a/gensrc/thrift/Partitions.thrift b/gensrc/thrift/Partitions.thrift index b14e36a3f628ee..95ff0eb9487af5 100644 --- a/gensrc/thrift/Partitions.thrift +++ b/gensrc/thrift/Partitions.thrift @@ -52,7 +52,13 @@ enum TPartitionType { HIVE_TABLE_SINK_UNPARTITIONED = 8, // used for merge partitioning: insert by partition columns, delete by row_id - MERGE_PARTITIONED = 9 + MERGE_PARTITIONED = 9, + + // connector-owned ownership function followed by writer assignment + EXTERNAL_TABLE_SINK_HASH_PARTITIONED = 10, + + // adaptive writer distribution without an ownership key + EXTERNAL_TABLE_SINK_UNPARTITIONED = 11 } enum TLocalPartitionType { @@ -194,6 +200,19 @@ struct TMergePartitionInfo { 6: optional i32 partition_spec_id } +enum TExternalTableSinkWriterAssignment { + IDENTITY = 0, + SKEWED = 1 +} + +// FE treats partition_function and its options as opaque connector-owned data. +// BE validates the named function before processing rows. +struct TExternalTableSinkHashPartitionInfo { + 1: required string partition_function + 2: optional map partition_function_options + 3: required TExternalTableSinkWriterAssignment writer_assignment +} + // Specification of how a single logical data stream is partitioned. // This leaves out the parameters that determine the physical partition (for hash // partitions, the number of partitions; for range partitions, the partitions' @@ -203,4 +222,5 @@ struct TDataPartition { 2: optional list partition_exprs 3: optional list partition_infos 4: optional TMergePartitionInfo merge_partition_info + 5: optional TExternalTableSinkHashPartitionInfo external_table_sink_hash_partition_info }