Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
133 changes: 133 additions & 0 deletions be/src/core/data_type_serde/data_type_variant_v2_serde.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand Down Expand Up @@ -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) {}
Expand Down Expand Up @@ -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());
});
Expand Down
38 changes: 36 additions & 2 deletions be/src/exec/operator/exchange_sink_operator.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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");
Expand Down Expand Up @@ -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));
Expand Down Expand Up @@ -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";
Expand All @@ -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
Expand Down Expand Up @@ -535,9 +567,11 @@ 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();
Expand Down
2 changes: 2 additions & 0 deletions be/src/exec/operator/exchange_sink_operator.h
Original file line number Diff line number Diff line change
Expand Up @@ -243,6 +243,8 @@ class ExchangeSinkOperatorX MOCK_REMOVE(final) : public DataSinkOperatorX<Exchan
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;

Expand Down
Original file line number Diff line number Diff line change
@@ -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
Original file line number Diff line number Diff line change
@@ -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 <gen_cpp/Partitions_types.h>

#include <memory>
#include <vector>

#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<TExpr>& partition_exprs,
std::unique_ptr<PartitionFunction>* partition_function);

} // namespace doris
Loading
Loading