eldenmoon commented on code in PR #65561: URL: https://github.com/apache/doris/pull/65561#discussion_r3613346971
########## be/src/exprs/function/cast/variant_v2/cast_variant_v2_document.cpp: ########## @@ -0,0 +1,225 @@ +// 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 <utility> + +#include "core/assert_cast.h" +#include "core/column/column_string.h" +#include "core/column/variant_v2/column_variant_v2.h" +#include "core/custom_allocator.h" +#include "core/data_type/data_type_string.h" +#include "core/string_buffer.hpp" +#include "exprs/function/cast/variant_v2/cast_variant_v2_internal.h" +#include "exprs/function_context.h" +#include "runtime/runtime_state.h" +#include "util/jsonb_writer.h" +#include "util/variant/variant_json.h" +#include "util/variant/variant_jsonb.h" + +namespace doris::CastWrapper::variant_v2_internal { +namespace { + +VariantJsonFormatOptions json_options(FunctionContext* context) { + if (context == nullptr || context->state() == nullptr) { + return {}; + } + return {.timezone = &context->state()->timezone_obj()}; +} + +bool uses_concrete_string_cast(VariantValueRef value) { + if (value.basic_type() == VariantBasicType::SHORT_STRING) { + return true; + } + if (value.basic_type() != VariantBasicType::PRIMITIVE) { + return false; + } + switch (value.primitive_id()) { + 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::STRING: + case VariantPrimitiveId::TIMESTAMP_NANOS: + case VariantPrimitiveId::TIMESTAMP_NTZ_NANOS: + return true; + case VariantPrimitiveId::NULL_VALUE: + case VariantPrimitiveId::BINARY: + case VariantPrimitiveId::TIME_NTZ_MICROS: + case VariantPrimitiveId::UUID: + return false; + } + throw Exception(ErrorCode::CORRUPTION, "Unknown Variant primitive id"); +} + +const ColumnVariantV2& encoded_source(const ColumnVariantV2& source, ColumnPtr* owner) { + if (!source.is_typed()) { + return source; + } + Status status = clone_as_encoded(source, owner); + if (!status.ok()) { + throw Exception(status); + } + return assert_cast<const ColumnVariantV2&>(**owner); +} + +Status append_fallback_string(VariantValueRef value, bool forced_null, + const VariantJsonFormatOptions& options, VectorBufferWriter* writer, + ColumnUInt8* nulls) { + if (forced_null) { + nulls->insert_value(1); + writer->commit(); + return Status::OK(); + } + nulls->insert_value(0); + if (value.is_null()) { + writer->write("null", 4); + } else { + to_json(value, *writer, options); + } + writer->commit(); + return Status::OK(); +} + +} // namespace + +Status cast_jsonb_to_variant(const ColumnPtr& source, size_t rows, ForcedNulls forced_nulls, + ColumnPtr* output) { + const auto* strings = check_and_get_column<ColumnString>(source.get()); + if (strings == nullptr || strings->size() != rows || + (!forced_nulls.empty() && forced_nulls.size() != rows)) { + return Status::InvalidArgument("Invalid JSONB input shape for Variant V2 CAST"); + } + JsonbToVariantEncoder encoder(VariantBlockBuilder::ReserveHint {.rows = rows}); + for (size_t row = 0; row < rows; ++row) { + if (!forced_nulls.empty() && forced_nulls[row] != 0) { + encoder.add_null(); + } else { + encoder.add_jsonb(strings->get_data_at(row)); + } + } + VariantEncodedBlock block = encoder.finish_block(); + auto result = ColumnVariantV2::create(); + result->insert_encoded_block(block.view()); + *output = std::move(result); + return Status::OK(); +} + +Status cast_variant_to_string(FunctionContext* context, const ColumnVariantV2& source, size_t rows, Review Comment: 这里函数拆成单独的文件，cast_variant_to_string.cpp -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
