[feature](external) Add generic connector write framework prerequisites

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 <arrow/array/builder_binary.h>
+#include <arrow/array/builder_nested.h>
 
 #include <algorithm>
 #include <cstring>
@@ -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<uint8_t>(primitive_id));
+    }
+    throw Exception(ErrorCode::NOT_IMPLEMENTED_ERROR,
+                    "Binary Variant V2 Arrow encoding does not support unknown primitive id {}",
+                    static_cast<uint8_t>(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<const arrow::StructType&>(*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<value: binary, metadata: binary>, got {}",
+                struct_type.ToString());
+    }
+    auto* value_builder = dynamic_cast<arrow::BinaryBuilder*>(builder.field_builder(0));
+    auto* metadata_builder = dynamic_cast<arrow::BinaryBuilder*>(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<const uint8_t*>(value.value.data),
+                                              cast_set<int32_t, size_t, false>(value.value.size)));
+                require_variant_arrow_status(metadata_builder->Append(
+                        reinterpret_cast<const uint8_t*>(value.metadata.data),
+                        cast_set<int32_t, size_t, false>(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<arrow::LargeStringBuilder&>(*array_builder), first, last,
                                options);
         }
+        if (array_builder->type()->id() == arrow::Type::STRUCT) {
+            return write_binary_variant_arrow(column, null_map,
+                                              assert_cast<arrow::StructBuilder&>(*array_builder),
+                                              first, last);
+        }
         return Status::InvalidArgument("Unsupported arrow type for variant column: {}",
                                        array_builder->type()->name());
     });

be/src/exec/operator/exchange_sink_operator.cpp

@@ -28,13 +28,15 @@
 #include <random>
 #include <string>
 
+#include "agent/be_exec_version_manager.h"
 #include "common/status.h"
 #include "core/column/column_const.h"
 #include "exec/exchange/exchange_writer.h"
 #include "exec/exchange/local_exchange_sink_operator.h"
 #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<ExternalTableSinkHashPartitioner>(
+                _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);

be/src/exec/operator/exchange_sink_operator.h

@@ -200,6 +200,17 @@ class ExchangeSinkOperatorX MOCK_REMOVE(final) : public DataSinkOperatorX<Exchan
 
     Status sink_impl(RuntimeState* state, Block* in_block, bool eos) override;
 
+#ifdef BE_TEST
+    size_t update_writer_scaling_for_test(size_t block_bytes, size_t max_writer_count) {
+        return _update_writer_scaling(block_bytes, max_writer_count);
+    }
+
+    std::pair<size_t, size_t> 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<Exchan
     // or each ExchangeSinkLocalState can have its own sink buffer.
     std::shared_ptr<ExchangeSinkBuffer> _create_buffer(
             RuntimeState* state, const std::vector<InstanceLoId>& sender_ins_ids);
+    size_t _update_writer_scaling(size_t block_bytes, size_t max_writer_count);
     std::shared_ptr<ExchangeSinkBuffer> _sink_buffer = nullptr;
     RuntimeState* _state = nullptr;
 
     const std::vector<TExpr> _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<Exchan
 
     // for external table sink random partition
     // Control the number of channels according to the flow, thereby controlling the number of table sink writers.
-    size_t _data_processed = 0;
-    int _writer_count = 1;
+    AnnotatedMutex _writer_scaling_mutex;
+    size_t _data_processed GUARDED_BY(_writer_scaling_mutex) = 0;
+    size_t _writer_count GUARDED_BY(_writer_scaling_mutex) = 1;
     // If dest_is_merge is true, it indicates that the corresponding receiver is a VMERGING-EXCHANGE.
     // The receiver will sort the collected data, so the sender must ensure that the data sent is ordered.
     const bool _dest_is_merge;

be/src/exec/partitioner/external/external_partition_function_factory.cpp

@@ -0,0 +1,45 @@
+// 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 "common/status.h"
+
+namespace doris {
+
+Status create_external_partition_function(const TExternalTableSinkHashPartitionInfo& partition_info,
+                                          PartitionerBase::HashValType logical_partition_count,
+                                          ShuffleHashMethod hash_method,
+                                          const std::vector<TExpr>& partition_exprs,
+                                          std::unique_ptr<PartitionFunction>* 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<HashPartitionFunction>(logical_partition_count, hash_method);
+    RETURN_IF_ERROR(function->init(partition_exprs));
+    *partition_function = std::move(function);
+    return Status::OK();
+}
+
+} // namespace doris