| // 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 <fmt/format.h> |
| |
| #include <cstddef> |
| #include <cstring> |
| #include <string> |
| #include <string_view> |
| #include <vector> |
| |
| #include "common/status.h" |
| #include "core/assert_cast.h" |
| #include "core/block/block.h" |
| #include "core/block/column_numbers.h" |
| #include "core/column/column.h" |
| #include "core/column/column_array.h" |
| #include "core/column/column_const.h" |
| #include "core/column/column_nullable.h" |
| #include "core/column/column_string.h" |
| #include "core/column/column_vector.h" |
| #include "core/data_type/data_type_array.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" |
| #include "core/memcpy_small.h" |
| #include "core/string_ref.h" |
| #include "exec/common/stringop_substring.h" |
| #include "exec/common/template_helpers.hpp" |
| #include "exec/common/util.hpp" |
| #include "exprs/function/function.h" |
| #include "exprs/function/function_helpers.h" |
| #include "exprs/function_context.h" |
| #include "util/simd/vstring_function.h" |
| |
| namespace doris { |
| #include "common/compile_check_avoid_begin.h" |
| |
| class FunctionStringConcat : public IFunction { |
| public: |
| struct ConcatState { |
| bool use_state = false; |
| std::string tail; |
| }; |
| |
| static constexpr auto name = "concat"; |
| static FunctionPtr create() { return std::make_shared<FunctionStringConcat>(); } |
| String get_name() const override { return name; } |
| size_t get_number_of_arguments() const override { return 0; } |
| bool is_variadic() const override { return true; } |
| |
| DataTypePtr get_return_type_impl(const DataTypes& arguments) const override { |
| return std::make_shared<DataTypeString>(); |
| } |
| |
| Status open(FunctionContext* context, FunctionContext::FunctionStateScope scope) override { |
| if (scope == FunctionContext::THREAD_LOCAL) { |
| return Status::OK(); |
| } |
| std::shared_ptr<ConcatState> state = std::make_shared<ConcatState>(); |
| |
| context->set_function_state(scope, state); |
| |
| state->use_state = true; |
| |
| // Optimize function calls like this: |
| // concat(col, "123", "abc", "456") -> tail = "123abc456" |
| for (size_t i = 1; i < context->get_num_args(); i++) { |
| const auto* column_string = context->get_constant_col(i); |
| if (column_string == nullptr) { |
| state->use_state = false; |
| return IFunction::open(context, scope); |
| } |
| auto string_vale = column_string->column_ptr->get_data_at(0); |
| if (string_vale.data == nullptr) { |
| // For concat(col, null), it is handled by default_implementation_for_nulls |
| state->use_state = false; |
| return IFunction::open(context, scope); |
| } |
| |
| state->tail.append(string_vale.begin(), string_vale.size); |
| } |
| |
| // The reserve is used here to allow the usage of memcpy_small_allow_read_write_overflow15 below. |
| state->tail.reserve(state->tail.size() + 16); |
| |
| return IFunction::open(context, scope); |
| } |
| |
| Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments, |
| uint32_t result, size_t input_rows_count) const override { |
| DCHECK_GE(arguments.size(), 1); |
| |
| if (arguments.size() == 1) { |
| block.get_by_position(result).column = block.get_by_position(arguments[0]).column; |
| return Status::OK(); |
| } |
| auto* concat_state = reinterpret_cast<ConcatState*>( |
| context->get_function_state(FunctionContext::FRAGMENT_LOCAL)); |
| if (!concat_state) { |
| return Status::RuntimeError("funciton context for function '{}' must have ConcatState;", |
| get_name()); |
| } |
| if (concat_state->use_state) { |
| const auto& [col, is_const] = |
| unpack_if_const(block.get_by_position(arguments[0]).column); |
| const auto* col_str = assert_cast<const ColumnString*>(col.get()); |
| if (is_const) { |
| return execute_const<true>(concat_state, block, col_str, result, input_rows_count); |
| } else { |
| return execute_const<false>(concat_state, block, col_str, result, input_rows_count); |
| } |
| |
| } else { |
| return execute_vecotr(block, arguments, result, input_rows_count); |
| } |
| } |
| |
| Status execute_vecotr(Block& block, const ColumnNumbers& arguments, uint32_t result, |
| size_t input_rows_count) const { |
| int argument_size = arguments.size(); |
| std::vector<ColumnPtr> argument_columns(argument_size); |
| |
| std::vector<const ColumnString::Offsets*> offsets_list(argument_size); |
| std::vector<const ColumnString::Chars*> chars_list(argument_size); |
| std::vector<bool> is_const_args(argument_size); |
| |
| for (int i = 0; i < argument_size; ++i) { |
| const auto& [col, is_const] = |
| unpack_if_const(block.get_by_position(arguments[i]).column); |
| |
| const auto* col_str = assert_cast<const ColumnString*>(col.get()); |
| offsets_list[i] = &col_str->get_offsets(); |
| chars_list[i] = &col_str->get_chars(); |
| is_const_args[i] = is_const; |
| } |
| |
| auto res = ColumnString::create(); |
| auto& res_data = res->get_chars(); |
| auto& res_offset = res->get_offsets(); |
| |
| res_offset.resize(input_rows_count); |
| size_t res_reserve_size = 0; |
| for (size_t i = 0; i < argument_size; ++i) { |
| if (is_const_args[i]) { |
| res_reserve_size += (*offsets_list[i])[0] * input_rows_count; |
| } else { |
| res_reserve_size += (*offsets_list[i])[input_rows_count - 1]; |
| } |
| } |
| |
| ColumnString::check_chars_length(res_reserve_size, 0); |
| |
| res_data.resize(res_reserve_size); |
| |
| auto* data = res_data.data(); |
| size_t dst_offset = 0; |
| |
| for (size_t i = 0; i < input_rows_count; ++i) { |
| for (size_t j = 0; j < argument_size; ++j) { |
| const auto& current_offsets = *offsets_list[j]; |
| const auto& current_chars = *chars_list[j]; |
| auto idx = index_check_const(i, is_const_args[j]); |
| const auto size = current_offsets[idx] - current_offsets[idx - 1]; |
| if (size > 0) { |
| memcpy_small_allow_read_write_overflow15( |
| data + dst_offset, current_chars.data() + current_offsets[idx - 1], |
| size); |
| dst_offset += size; |
| } |
| } |
| res_offset[i] = dst_offset; |
| } |
| |
| block.get_by_position(result).column = std::move(res); |
| return Status::OK(); |
| } |
| |
| template <bool is_const> |
| Status execute_const(ConcatState* concat_state, Block& block, const ColumnString* col_str, |
| uint32_t result, size_t input_rows_count) const { |
| // using tail optimize |
| |
| auto res = ColumnString::create(); |
| auto& res_data = res->get_chars(); |
| auto& res_offset = res->get_offsets(); |
| res_offset.resize(input_rows_count); |
| |
| size_t res_reserve_size = 0; |
| if constexpr (is_const) { |
| res_reserve_size = col_str->get_offsets()[0] * input_rows_count; |
| } else { |
| res_reserve_size = col_str->get_offsets()[input_rows_count - 1]; |
| } |
| res_reserve_size += concat_state->tail.size() * input_rows_count; |
| |
| ColumnString::check_chars_length(res_reserve_size, 0); |
| res_data.resize(res_reserve_size); |
| |
| const auto& tail = concat_state->tail; |
| auto* data = res_data.data(); |
| size_t dst_offset = 0; |
| |
| for (size_t i = 0; i < input_rows_count; ++i) { |
| const auto idx = index_check_const<is_const>(i); |
| StringRef str_val = col_str->get_data_at(idx); |
| // copy column |
| memcpy_small_allow_read_write_overflow15(data + dst_offset, str_val.data, str_val.size); |
| dst_offset += str_val.size; |
| // copy tail |
| memcpy_small_allow_read_write_overflow15(data + dst_offset, tail.data(), tail.size()); |
| dst_offset += tail.size(); |
| res_offset[i] = dst_offset; |
| } |
| block.get_by_position(result).column = std::move(res); |
| return Status::OK(); |
| } |
| }; |
| |
| class FunctionStringElt : public IFunction { |
| public: |
| static constexpr auto name = "elt"; |
| static FunctionPtr create() { return std::make_shared<FunctionStringElt>(); } |
| String get_name() const override { return name; } |
| size_t get_number_of_arguments() const override { return 0; } |
| bool is_variadic() const override { return true; } |
| |
| DataTypePtr get_return_type_impl(const DataTypes& arguments) const override { |
| return make_nullable(std::make_shared<DataTypeString>()); |
| } |
| bool use_default_implementation_for_nulls() const override { return false; } |
| |
| Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments, |
| uint32_t result, size_t input_rows_count) const override { |
| int arguent_size = arguments.size(); |
| int num_children = arguent_size - 1; |
| auto res = ColumnString::create(); |
| |
| if (auto const_column = check_and_get_column<ColumnConst>( |
| *block.get_by_position(arguments[0]).column)) { |
| auto data = const_column->get_data_at(0); |
| // return NULL, pos is null or pos < 0 or pos > num_children |
| auto is_null = data.data == nullptr; |
| auto pos = is_null ? 0 : *(Int32*)data.data; |
| is_null = pos <= 0 || pos > num_children; |
| |
| auto null_map = ColumnUInt8::create(input_rows_count, is_null); |
| if (is_null) { |
| res->insert_many_defaults(input_rows_count); |
| } else { |
| auto& target_column = block.get_by_position(arguments[pos]).column; |
| if (auto target_const_column = check_and_get_column<ColumnConst>(*target_column)) { |
| auto target_data = target_const_column->get_data_at(0); |
| // return NULL, no target data |
| if (target_data.data == nullptr) { |
| null_map = ColumnUInt8::create(input_rows_count, true); |
| res->insert_many_defaults(input_rows_count); |
| } else { |
| res->insert_data_repeatedly(target_data.data, target_data.size, |
| input_rows_count); |
| } |
| } else if (auto target_nullable_column = |
| check_and_get_column<ColumnNullable>(*target_column)) { |
| auto& target_null_map = target_nullable_column->get_null_map_data(); |
| VectorizedUtils::update_null_map(null_map->get_data(), target_null_map); |
| |
| auto& target_str_column = assert_cast<const ColumnString&>( |
| target_nullable_column->get_nested_column()); |
| res->get_chars().assign(target_str_column.get_chars().begin(), |
| target_str_column.get_chars().end()); |
| res->get_offsets().assign(target_str_column.get_offsets().begin(), |
| target_str_column.get_offsets().end()); |
| } else { |
| auto& target_str_column = assert_cast<const ColumnString&>(*target_column); |
| res->get_chars().assign(target_str_column.get_chars().begin(), |
| target_str_column.get_chars().end()); |
| res->get_offsets().assign(target_str_column.get_offsets().begin(), |
| target_str_column.get_offsets().end()); |
| } |
| } |
| block.get_by_position(result).column = |
| ColumnNullable::create(std::move(res), std::move(null_map)); |
| } else if (auto pos_null_column = check_and_get_column<ColumnNullable>( |
| *block.get_by_position(arguments[0]).column)) { |
| auto& pos_column = |
| assert_cast<const ColumnInt32&>(pos_null_column->get_nested_column()); |
| auto& pos_null_map = pos_null_column->get_null_map_data(); |
| auto null_map = ColumnUInt8::create(input_rows_count, false); |
| auto& res_null_map = null_map->get_data(); |
| |
| for (size_t i = 0; i < input_rows_count; ++i) { |
| auto pos = pos_column.get_element(i); |
| res_null_map[i] = |
| pos_null_map[i] || pos <= 0 || pos > num_children || |
| block.get_by_position(arguments[pos]).column->get_data_at(i).data == |
| nullptr; |
| if (res_null_map[i]) { |
| res->insert_default(); |
| } else { |
| auto insert_data = block.get_by_position(arguments[pos]).column->get_data_at(i); |
| res->insert_data(insert_data.data, insert_data.size); |
| } |
| } |
| block.get_by_position(result).column = |
| ColumnNullable::create(std::move(res), std::move(null_map)); |
| } else { |
| auto& pos_column = |
| assert_cast<const ColumnInt32&>(*block.get_by_position(arguments[0]).column); |
| auto null_map = ColumnUInt8::create(input_rows_count, false); |
| auto& res_null_map = null_map->get_data(); |
| |
| for (size_t i = 0; i < input_rows_count; ++i) { |
| auto pos = pos_column.get_element(i); |
| res_null_map[i] = |
| pos <= 0 || pos > num_children || |
| block.get_by_position(arguments[pos]).column->get_data_at(i).data == |
| nullptr; |
| if (res_null_map[i]) { |
| res->insert_default(); |
| } else { |
| auto insert_data = block.get_by_position(arguments[pos]).column->get_data_at(i); |
| res->insert_data(insert_data.data, insert_data.size); |
| } |
| } |
| block.get_by_position(result).column = |
| ColumnNullable::create(std::move(res), std::move(null_map)); |
| } |
| return Status::OK(); |
| } |
| }; |
| |
| // concat_ws (string,string....) or (string, Array) |
| // TODO: avoid use fmtlib |
| class FunctionStringConcatWs : public IFunction { |
| public: |
| using Chars = ColumnString::Chars; |
| using Offsets = ColumnString::Offsets; |
| |
| static constexpr auto name = "concat_ws"; |
| static FunctionPtr create() { return std::make_shared<FunctionStringConcatWs>(); } |
| String get_name() const override { return name; } |
| size_t get_number_of_arguments() const override { return 0; } |
| bool is_variadic() const override { return true; } |
| |
| DataTypePtr get_return_type_impl(const DataTypes& arguments) const override { |
| const IDataType* first_type = arguments[0].get(); |
| if (first_type->is_nullable()) { |
| return make_nullable(std::make_shared<DataTypeString>()); |
| } else { |
| return std::make_shared<DataTypeString>(); |
| } |
| } |
| bool use_default_implementation_for_nulls() const override { return false; } |
| |
| Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments, |
| uint32_t result, size_t input_rows_count) const override { |
| DCHECK_GE(arguments.size(), 2); |
| auto null_map = ColumnUInt8::create(input_rows_count, 0); |
| // we create a zero column to simply implement |
| auto const_null_map = ColumnUInt8::create(input_rows_count, 0); |
| auto res = ColumnString::create(); |
| bool is_null_type = block.get_by_position(arguments[0]).type.get()->is_nullable(); |
| size_t argument_size = arguments.size(); |
| std::vector<const Offsets*> offsets_list(argument_size); |
| std::vector<const Chars*> chars_list(argument_size); |
| std::vector<const ColumnUInt8::Container*> null_list(argument_size); |
| |
| std::vector<ColumnPtr> argument_columns(argument_size); |
| std::vector<ColumnPtr> argument_null_columns(argument_size); |
| |
| for (size_t i = 0; i < argument_size; ++i) { |
| argument_columns[i] = |
| block.get_by_position(arguments[i]).column->convert_to_full_column_if_const(); |
| if (const auto* nullable = |
| check_and_get_column<const ColumnNullable>(*argument_columns[i])) { |
| // Danger: Here must dispose the null map data first! Because |
| // argument_columns[i]=nullable->get_nested_column_ptr(); will release the mem |
| // of column nullable mem of null map |
| null_list[i] = &nullable->get_null_map_data(); |
| argument_null_columns[i] = nullable->get_null_map_column_ptr(); |
| argument_columns[i] = nullable->get_nested_column_ptr(); |
| } else { |
| null_list[i] = &const_null_map->get_data(); |
| } |
| |
| if (is_column<ColumnArray>(argument_columns[i].get())) { |
| continue; |
| } |
| |
| const auto* col_str = assert_cast<const ColumnString*>(argument_columns[i].get()); |
| offsets_list[i] = &col_str->get_offsets(); |
| chars_list[i] = &col_str->get_chars(); |
| } |
| |
| auto& res_data = res->get_chars(); |
| auto& res_offset = res->get_offsets(); |
| res_offset.resize(input_rows_count); |
| |
| VectorizedUtils::update_null_map(null_map->get_data(), *null_list[0]); |
| fmt::memory_buffer buffer; |
| std::vector<std::string_view> views; |
| |
| if (is_column<ColumnArray>(argument_columns[1].get())) { |
| // Determine if the nested type of the array is String |
| const auto& array_column = reinterpret_cast<const ColumnArray&>(*argument_columns[1]); |
| if (!array_column.get_data().is_column_string()) { |
| return Status::NotSupported( |
| fmt::format("unsupported nested array of type {} for function {}", |
| is_column_nullable(array_column.get_data()) |
| ? array_column.get_data().get_name() |
| : array_column.get_data().get_name(), |
| get_name())); |
| } |
| // Concat string in array |
| _execute_array(input_rows_count, array_column, buffer, views, offsets_list, chars_list, |
| null_list, res_data, res_offset); |
| |
| } else { |
| // Concat string |
| _execute_string(input_rows_count, argument_size, buffer, views, offsets_list, |
| chars_list, null_list, res_data, res_offset); |
| } |
| if (is_null_type) { |
| block.get_by_position(result).column = |
| ColumnNullable::create(std::move(res), std::move(null_map)); |
| } else { |
| block.get_by_position(result).column = std::move(res); |
| } |
| return Status::OK(); |
| } |
| |
| private: |
| void _execute_array(const size_t& input_rows_count, const ColumnArray& array_column, |
| fmt::memory_buffer& buffer, std::vector<std::string_view>& views, |
| const std::vector<const Offsets*>& offsets_list, |
| const std::vector<const Chars*>& chars_list, |
| const std::vector<const ColumnUInt8::Container*>& null_list, |
| Chars& res_data, Offsets& res_offset) const { |
| // Get array nested column |
| const UInt8* array_nested_null_map = nullptr; |
| ColumnPtr array_nested_column = nullptr; |
| |
| if (is_column_nullable(array_column.get_data())) { |
| const auto& array_nested_null_column = |
| reinterpret_cast<const ColumnNullable&>(array_column.get_data()); |
| // String's null map in array |
| array_nested_null_map = |
| array_nested_null_column.get_null_map_column().get_data().data(); |
| array_nested_column = array_nested_null_column.get_nested_column_ptr(); |
| } else { |
| array_nested_column = array_column.get_data_ptr(); |
| } |
| |
| const auto& string_column = reinterpret_cast<const ColumnString&>(*array_nested_column); |
| const Chars& string_src_chars = string_column.get_chars(); |
| const auto& src_string_offsets = string_column.get_offsets(); |
| const auto& src_array_offsets = array_column.get_offsets(); |
| size_t current_src_array_offset = 0; |
| |
| // Concat string in array |
| for (size_t i = 0; i < input_rows_count; ++i) { |
| auto& sep_offsets = *offsets_list[0]; |
| auto& sep_chars = *chars_list[0]; |
| auto& sep_nullmap = *null_list[0]; |
| |
| if (sep_nullmap[i]) { |
| res_offset[i] = res_data.size(); |
| current_src_array_offset += src_array_offsets[i] - src_array_offsets[i - 1]; |
| continue; |
| } |
| |
| int sep_size = sep_offsets[i] - sep_offsets[i - 1]; |
| const char* sep_data = reinterpret_cast<const char*>(&sep_chars[sep_offsets[i - 1]]); |
| |
| std::string_view sep(sep_data, sep_size); |
| buffer.clear(); |
| views.clear(); |
| |
| for (auto next_src_array_offset = src_array_offsets[i]; |
| current_src_array_offset < next_src_array_offset; ++current_src_array_offset) { |
| const auto current_src_string_offset = |
| current_src_array_offset ? src_string_offsets[current_src_array_offset - 1] |
| : 0; |
| size_t bytes_to_copy = |
| src_string_offsets[current_src_array_offset] - current_src_string_offset; |
| const char* ptr = |
| reinterpret_cast<const char*>(&string_src_chars[current_src_string_offset]); |
| |
| if (array_nested_null_map == nullptr || |
| !array_nested_null_map[current_src_array_offset]) { |
| views.emplace_back(ptr, bytes_to_copy); |
| } |
| } |
| |
| fmt::format_to(buffer, "{}", fmt::join(views, sep)); |
| |
| StringOP::push_value_string(std::string_view(buffer.data(), buffer.size()), i, res_data, |
| res_offset); |
| } |
| } |
| |
| void _execute_string(const size_t& input_rows_count, const size_t& argument_size, |
| fmt::memory_buffer& buffer, std::vector<std::string_view>& views, |
| const std::vector<const Offsets*>& offsets_list, |
| const std::vector<const Chars*>& chars_list, |
| const std::vector<const ColumnUInt8::Container*>& null_list, |
| Chars& res_data, Offsets& res_offset) const { |
| // Concat string |
| for (size_t i = 0; i < input_rows_count; ++i) { |
| auto& sep_offsets = *offsets_list[0]; |
| auto& sep_chars = *chars_list[0]; |
| auto& sep_nullmap = *null_list[0]; |
| if (sep_nullmap[i]) { |
| res_offset[i] = res_data.size(); |
| continue; |
| } |
| |
| int sep_size = sep_offsets[i] - sep_offsets[i - 1]; |
| const char* sep_data = reinterpret_cast<const char*>(&sep_chars[sep_offsets[i - 1]]); |
| |
| std::string_view sep(sep_data, sep_size); |
| buffer.clear(); |
| views.clear(); |
| for (size_t j = 1; j < argument_size; ++j) { |
| auto& current_offsets = *offsets_list[j]; |
| auto& current_chars = *chars_list[j]; |
| auto& current_nullmap = *null_list[j]; |
| int size = current_offsets[i] - current_offsets[i - 1]; |
| const char* ptr = |
| reinterpret_cast<const char*>(¤t_chars[current_offsets[i - 1]]); |
| if (!current_nullmap[i]) { |
| views.emplace_back(ptr, size); |
| } |
| } |
| fmt::format_to(buffer, "{}", fmt::join(views, sep)); |
| StringOP::push_value_string(std::string_view(buffer.data(), buffer.size()), i, res_data, |
| res_offset); |
| } |
| } |
| }; |
| |
| class FunctionStringRepeat : public IFunction { |
| public: |
| static constexpr auto name = "repeat"; |
| static FunctionPtr create() { return std::make_shared<FunctionStringRepeat>(); } |
| String get_name() const override { return name; } |
| size_t get_number_of_arguments() const override { return 2; } |
| // should set NULL value of nested data to default, |
| // as iff it's not inited and invalid, the repeat result of length is so large cause overflow |
| bool need_replace_null_data_to_default() const override { return true; } |
| |
| DataTypePtr get_return_type_impl(const DataTypes& arguments) const override { |
| return make_nullable(std::make_shared<DataTypeString>()); |
| } |
| Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments, |
| uint32_t result, size_t input_rows_count) const override { |
| DCHECK_EQ(arguments.size(), 2); |
| auto res = ColumnString::create(); |
| auto null_map = ColumnUInt8::create(); |
| |
| ColumnPtr argument_ptr[2]; |
| argument_ptr[0] = |
| block.get_by_position(arguments[0]).column->convert_to_full_column_if_const(); |
| argument_ptr[1] = block.get_by_position(arguments[1]).column; |
| |
| if (const auto* col1 = check_and_get_column<ColumnString>(*argument_ptr[0])) { |
| if (const auto* col2 = check_and_get_column<ColumnInt32>(*argument_ptr[1])) { |
| RETURN_IF_ERROR(vector_vector(col1->get_chars(), col1->get_offsets(), |
| col2->get_data(), res->get_chars(), |
| res->get_offsets(), null_map->get_data())); |
| block.replace_by_position( |
| result, ColumnNullable::create(std::move(res), std::move(null_map))); |
| return Status::OK(); |
| } else if (const auto* col2_const = |
| check_and_get_column<ColumnConst>(*argument_ptr[1])) { |
| DCHECK(check_and_get_column<ColumnInt32>(col2_const->get_data_column())); |
| int repeat = col2_const->get_int(0); |
| if (repeat <= 0) { |
| null_map->get_data().resize_fill(input_rows_count, 0); |
| res->insert_many_defaults(input_rows_count); |
| } else { |
| vector_const(col1->get_chars(), col1->get_offsets(), repeat, res->get_chars(), |
| res->get_offsets(), null_map->get_data()); |
| } |
| block.replace_by_position( |
| result, ColumnNullable::create(std::move(res), std::move(null_map))); |
| return Status::OK(); |
| } |
| } |
| |
| return Status::RuntimeError("repeat function get error param: {}, {}", |
| argument_ptr[0]->get_name(), argument_ptr[1]->get_name()); |
| } |
| |
| Status vector_vector(const ColumnString::Chars& data, const ColumnString::Offsets& offsets, |
| const ColumnInt32::Container& repeats, ColumnString::Chars& res_data, |
| ColumnString::Offsets& res_offsets, |
| ColumnUInt8::Container& null_map) const { |
| size_t input_row_size = offsets.size(); |
| |
| fmt::memory_buffer buffer; |
| res_offsets.resize(input_row_size); |
| null_map.resize_fill(input_row_size, 0); |
| for (ssize_t i = 0; i < input_row_size; ++i) { |
| buffer.clear(); |
| const char* raw_str = reinterpret_cast<const char*>(&data[offsets[i - 1]]); |
| size_t size = offsets[i] - offsets[i - 1]; |
| int repeat = repeats[i]; |
| if (repeat <= 0) { |
| StringOP::push_empty_string(i, res_data, res_offsets); |
| } else { |
| ColumnString::check_chars_length(repeat * size + res_data.size(), 0); |
| for (int j = 0; j < repeat; ++j) { |
| buffer.append(raw_str, raw_str + size); |
| } |
| StringOP::push_value_string(std::string_view(buffer.data(), buffer.size()), i, |
| res_data, res_offsets); |
| } |
| } |
| return Status::OK(); |
| } |
| |
| // TODO: 1. use pmr::vector<char> replace fmt_buffer may speed up the code |
| // 2. abstract the `vector_vector` and `vector_const` |
| // 3. rethink we should use `DEFAULT_MAX_STRING_SIZE` to bigger here |
| void vector_const(const ColumnString::Chars& data, const ColumnString::Offsets& offsets, |
| int repeat, ColumnString::Chars& res_data, ColumnString::Offsets& res_offsets, |
| ColumnUInt8::Container& null_map) const { |
| size_t input_row_size = offsets.size(); |
| |
| fmt::memory_buffer buffer; |
| res_offsets.resize(input_row_size); |
| null_map.resize_fill(input_row_size, 0); |
| for (ssize_t i = 0; i < input_row_size; ++i) { |
| buffer.clear(); |
| const char* raw_str = reinterpret_cast<const char*>(&data[offsets[i - 1]]); |
| size_t size = offsets[i] - offsets[i - 1]; |
| ColumnString::check_chars_length(repeat * size + res_data.size(), 0); |
| |
| for (int j = 0; j < repeat; ++j) { |
| buffer.append(raw_str, raw_str + size); |
| } |
| StringOP::push_value_string(std::string_view(buffer.data(), buffer.size()), i, res_data, |
| res_offsets); |
| } |
| } |
| }; |
| |
| /// PaddingChars pre-processes the pad string for efficient padding. |
| /// When is_utf8=false, character count equals byte count — no UTF-8 decoding needed. |
| /// When is_utf8=true, we build a byte-offset table for code points. |
| /// In both cases, the pad string is pre-expanded (doubled) until it has >= 16 characters, |
| /// so that each memcpy in append_to copies at least 16 bytes at a time. |
| template <bool is_utf8> |
| struct PaddingChars { |
| std::string pad_string; |
| /// utf8_byte_offsets[i] = byte offset of i-th code point in pad_string. |
| /// utf8_byte_offsets has (num_chars + 1) entries, with [0]=0 and [num_chars]=pad_string.size(). |
| std::vector<size_t> utf8_byte_offsets; |
| |
| explicit PaddingChars(const uint8_t* data, size_t len) |
| : pad_string(reinterpret_cast<const char*>(data), len) { |
| init(); |
| } |
| |
| size_t num_chars() const { |
| if constexpr (is_utf8) { |
| return utf8_byte_offsets.size() - 1; |
| } else { |
| return pad_string.size(); |
| } |
| } |
| |
| size_t chars_to_bytes(size_t n) const { |
| if constexpr (is_utf8) { |
| return utf8_byte_offsets[n]; |
| } else { |
| return n; |
| } |
| } |
| |
| /// Append `num_chars_to_pad` padding characters to dst, return bytes written. |
| size_t append_to(uint8_t* dst, size_t num_chars_to_pad) const { |
| if (num_chars_to_pad == 0) { |
| return 0; |
| } |
| const auto* src = reinterpret_cast<const uint8_t*>(pad_string.data()); |
| const size_t step = num_chars(); |
| uint8_t* dst_start = dst; |
| while (num_chars_to_pad > step) { |
| size_t bytes = chars_to_bytes(step); |
| memcpy(dst, src, bytes); |
| dst += bytes; |
| num_chars_to_pad -= step; |
| } |
| size_t bytes = chars_to_bytes(num_chars_to_pad); |
| memcpy(dst, src, bytes); |
| dst += bytes; |
| return dst - dst_start; |
| } |
| |
| private: |
| void init() { |
| if (pad_string.empty()) { |
| return; |
| } |
| |
| if constexpr (is_utf8) { |
| // Build byte-offset table for each code point. |
| size_t offset = 0; |
| utf8_byte_offsets.reserve(pad_string.size() + 1); |
| while (offset < pad_string.size()) { |
| utf8_byte_offsets.push_back(offset); |
| offset += get_utf8_byte_length(static_cast<uint8_t>(pad_string[offset])); |
| offset = std::min(offset, pad_string.size()); |
| } |
| utf8_byte_offsets.push_back(pad_string.size()); |
| } |
| |
| // Pre-expand pad_string until it has >= 16 characters. |
| // This ensures append_to() copies at least 16 bytes per iteration. |
| while (num_chars() < 16) { |
| if constexpr (is_utf8) { |
| size_t old_count = utf8_byte_offsets.size(); |
| size_t base = utf8_byte_offsets.back(); |
| for (size_t i = 1; i < old_count; ++i) { |
| utf8_byte_offsets.push_back(utf8_byte_offsets[i] + base); |
| } |
| } |
| pad_string += pad_string; |
| } |
| } |
| }; |
| |
| template <typename Impl> |
| class FunctionStringPad : public IFunction { |
| public: |
| static constexpr auto name = Impl::name; |
| static FunctionPtr create() { return std::make_shared<FunctionStringPad>(); } |
| String get_name() const override { return name; } |
| size_t get_number_of_arguments() const override { return 3; } |
| |
| DataTypePtr get_return_type_impl(const DataTypes& arguments) const override { |
| return make_nullable(std::make_shared<DataTypeString>()); |
| } |
| |
| Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments, |
| uint32_t result, size_t input_rows_count) const override { |
| DCHECK_GE(arguments.size(), 3); |
| auto null_map = ColumnUInt8::create(input_rows_count, 0); |
| auto res = ColumnString::create(); |
| |
| ColumnPtr col[3]; |
| bool col_const[3]; |
| for (size_t i = 0; i < 3; ++i) { |
| std::tie(col[i], col_const[i]) = |
| unpack_if_const(block.get_by_position(arguments[i]).column); |
| } |
| auto& null_map_data = null_map->get_data(); |
| auto& res_offsets = res->get_offsets(); |
| auto& res_chars = res->get_chars(); |
| res_offsets.resize(input_rows_count); |
| |
| const auto* strcol = assert_cast<const ColumnString*>(col[0].get()); |
| const auto* col_len = assert_cast<const ColumnInt32*>(col[1].get()); |
| const auto& col_len_data = col_len->get_data(); |
| |
| const auto* padcol = assert_cast<const ColumnString*>(col[2].get()); |
| |
| if (col_const[1] && col_const[2]) { |
| auto pad = padcol->get_data_at(0); |
| const bool pad_all_ascii = |
| simd::VStringFunctions::is_ascii({pad.data, static_cast<size_t>(pad.size)}); |
| const bool all_ascii = pad_all_ascii && strcol->is_ascii(); |
| std::visit( |
| [&](auto str_const) { |
| if (all_ascii) { |
| execute_const_len_const_pad<true, str_const>( |
| *strcol, col_len_data, *padcol, res_offsets, res_chars, |
| null_map_data, input_rows_count); |
| } else { |
| execute_const_len_const_pad<false, str_const>( |
| *strcol, col_len_data, *padcol, res_offsets, res_chars, |
| null_map_data, input_rows_count); |
| } |
| }, |
| make_bool_variant(col_const[0])); |
| } else { |
| std::visit( |
| [&](auto str_const) { |
| execute_general<str_const>(*strcol, col_len_data, col_const[1], *padcol, |
| col_const[2], res_offsets, res_chars, |
| null_map_data, input_rows_count); |
| }, |
| make_bool_variant(col_const[0])); |
| } |
| |
| block.get_by_position(result).column = |
| ColumnNullable::create(std::move(res), std::move(null_map)); |
| return Status::OK(); |
| } |
| |
| private: |
| template <bool is_utf8> |
| static size_t get_char_length(const uint8_t* str_data, size_t str_byte_len) { |
| if constexpr (is_utf8) { |
| return simd::VStringFunctions::get_char_len(reinterpret_cast<const char*>(str_data), |
| str_byte_len); |
| } |
| return str_byte_len; |
| } |
| |
| template <bool is_utf8> |
| static size_t get_truncated_byte_length(const uint8_t* str_data, size_t str_byte_len, |
| size_t str_char_len, size_t target_len) { |
| if constexpr (!is_utf8) { |
| return target_len; |
| } |
| if (str_char_len == target_len) { |
| return str_byte_len; |
| } |
| auto [byte_len, _] = simd::VStringFunctions::iterate_utf8_with_limit_length( |
| reinterpret_cast<const char*>(str_data), |
| reinterpret_cast<const char*>(str_data) + str_byte_len, target_len); |
| return byte_len; |
| } |
| |
| static void ensure_capacity(ColumnString::Chars& res_chars, size_t needed, size_t row) { |
| if (needed <= res_chars.size()) { |
| return; |
| } |
| ColumnString::check_chars_length(needed, row); |
| res_chars.resize(std::max(needed, res_chars.size() * 3 / 2)); |
| } |
| |
| template <bool is_utf8> |
| static size_t estimate_const_output_bytes(const ColumnString::Chars& strcol_chars, |
| int target_len, size_t input_rows_count, |
| const PaddingChars<is_utf8>* padding) { |
| if (target_len <= 0) { |
| return 0; |
| } |
| if constexpr (!is_utf8) { |
| return static_cast<size_t>(target_len) * input_rows_count; |
| } |
| if (padding != nullptr && padding->num_chars() > 0) { |
| size_t pad_bytes_per_char = |
| (padding->pad_string.size() + padding->num_chars() - 1) / padding->num_chars(); |
| return strcol_chars.size() + |
| static_cast<size_t>(target_len) * pad_bytes_per_char * input_rows_count; |
| } |
| return strcol_chars.size(); |
| } |
| |
| template <bool is_utf8> |
| static void append_result_row(const uint8_t* str_data, size_t str_byte_len, int target_len, |
| const PaddingChars<is_utf8>* padding, |
| ColumnString::Chars& res_chars, |
| ColumnString::Offsets& res_offsets, |
| ColumnUInt8::Container& null_map_data, size_t row, |
| size_t& dst_offset) { |
| if (target_len < 0) { |
| null_map_data[row] = true; |
| res_offsets[row] = dst_offset; |
| return; |
| } |
| |
| const size_t str_char_len = get_char_length<is_utf8>(str_data, str_byte_len); |
| const size_t target_char_len = static_cast<size_t>(target_len); |
| if (str_char_len >= target_char_len) { |
| const size_t truncated_byte_len = get_truncated_byte_length<is_utf8>( |
| str_data, str_byte_len, str_char_len, target_char_len); |
| const size_t needed = dst_offset + truncated_byte_len; |
| ensure_capacity(res_chars, needed, row); |
| memcpy(res_chars.data() + dst_offset, str_data, truncated_byte_len); |
| dst_offset += truncated_byte_len; |
| res_offsets[row] = dst_offset; |
| return; |
| } |
| |
| if (padding == nullptr || padding->num_chars() == 0) { |
| res_offsets[row] = dst_offset; |
| return; |
| } |
| |
| const size_t pad_char_count = target_char_len - str_char_len; |
| const size_t full_cycles = pad_char_count / padding->num_chars(); |
| const size_t remainder_chars = pad_char_count % padding->num_chars(); |
| const size_t pad_bytes = |
| full_cycles * padding->pad_string.size() + padding->chars_to_bytes(remainder_chars); |
| const size_t needed = dst_offset + str_byte_len + pad_bytes; |
| ensure_capacity(res_chars, needed, row); |
| |
| if constexpr (Impl::is_lpad) { |
| dst_offset += padding->append_to(res_chars.data() + dst_offset, pad_char_count); |
| memcpy(res_chars.data() + dst_offset, str_data, str_byte_len); |
| dst_offset += str_byte_len; |
| } else { |
| memcpy(res_chars.data() + dst_offset, str_data, str_byte_len); |
| dst_offset += str_byte_len; |
| dst_offset += padding->append_to(res_chars.data() + dst_offset, pad_char_count); |
| } |
| res_offsets[row] = dst_offset; |
| } |
| |
| template <bool all_ascii, bool str_const> |
| static void execute_const_len_const_pad(const ColumnString& strcol, |
| const ColumnInt32::Container& col_len_data, |
| const ColumnString& padcol, |
| ColumnString::Offsets& res_offsets, |
| ColumnString::Chars& res_chars, |
| ColumnUInt8::Container& null_map_data, |
| size_t input_rows_count) { |
| constexpr bool is_utf8 = !all_ascii; |
| using PadChars = PaddingChars<is_utf8>; |
| |
| const int target_len = col_len_data[0]; |
| std::optional<PadChars> padding; |
| const auto pad = padcol.get_data_at(0); |
| if (!pad.empty()) { |
| padding.emplace(reinterpret_cast<const uint8_t*>(pad.data), pad.size); |
| } |
| |
| const PadChars* padding_ptr = padding ? &*padding : nullptr; |
| const size_t estimated_total = estimate_const_output_bytes<is_utf8>( |
| strcol.get_chars(), target_len, input_rows_count, padding_ptr); |
| if (estimated_total > 0) { |
| ColumnString::check_chars_length(estimated_total, 0, input_rows_count); |
| } |
| res_chars.resize(estimated_total); |
| |
| size_t dst_offset = 0; |
| for (size_t i = 0; i < input_rows_count; ++i) { |
| auto str = strcol.get_data_at(index_check_const<str_const>(i)); |
| append_result_row<is_utf8>(reinterpret_cast<const uint8_t*>(str.data), str.size, |
| target_len, padding_ptr, res_chars, res_offsets, |
| null_map_data, i, dst_offset); |
| } |
| res_chars.resize(dst_offset); |
| } |
| |
| template <bool str_const> |
| static void execute_general(const ColumnString& strcol, |
| const ColumnInt32::Container& col_len_data, bool len_const, |
| const ColumnString& padcol, bool pad_const, |
| ColumnString::Offsets& res_offsets, ColumnString::Chars& res_chars, |
| ColumnUInt8::Container& null_map_data, size_t input_rows_count) { |
| using PadChars = PaddingChars<true>; |
| std::optional<PadChars> const_padding; |
| const PadChars* const_padding_ptr = nullptr; |
| if (pad_const) { |
| auto pad = padcol.get_data_at(0); |
| if (!pad.empty()) { |
| const_padding.emplace(reinterpret_cast<const uint8_t*>(pad.data), pad.size); |
| const_padding_ptr = &*const_padding; |
| } |
| } |
| |
| res_chars.resize(strcol.get_chars().size()); |
| size_t dst_offset = 0; |
| for (size_t i = 0; i < input_rows_count; ++i) { |
| auto str = strcol.get_data_at(index_check_const<str_const>(i)); |
| const int target_len = col_len_data[len_const ? 0 : i]; |
| |
| const PadChars* padding_ptr = const_padding_ptr; |
| std::optional<PadChars> row_padding; |
| if (!pad_const) { |
| auto pad = padcol.get_data_at(i); |
| if (!pad.empty()) { |
| row_padding.emplace(reinterpret_cast<const uint8_t*>(pad.data), pad.size); |
| padding_ptr = &*row_padding; |
| } else { |
| padding_ptr = nullptr; |
| } |
| } |
| |
| append_result_row<true>(reinterpret_cast<const uint8_t*>(str.data), str.size, |
| target_len, padding_ptr, res_chars, res_offsets, null_map_data, |
| i, dst_offset); |
| } |
| res_chars.resize(dst_offset); |
| } |
| }; |
| |
| #include "common/compile_check_avoid_end.h" |
| } // namespace doris |
