be/src/exprs/function/function_regexp.cpp

Source
// 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 <glog/logging.h>
#include <re2/re2.h>
#include <re2/stringpiece.h>
#include <stddef.h>

#include <boost/regex.hpp>
#include <memory>
#include <string>
#include <string_view>
#include <type_traits>
#include <utility>
#include <vector>

#include "common/status.h"
#include "core/block/block.h"
#include "core/block/column_numbers.h"
#include "core/block/column_with_type_and_name.h"
#include "core/column/column.h"
#include "core/column/column_array.h"
#include "core/column/column_const.h"
#include "core/column/column_execute_util.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.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/string_ref.h"
#include "core/types.h"
#include "exec/common/stringop_substring.h"
#include "exprs/aggregate/aggregate_function.h"
#include "exprs/function/function.h"
#include "exprs/function/simple_function_factory.h"
#include "exprs/function_context.h"
#include "exprs/string_functions.h"

namespace doris {

// Helper structure to hold either RE2 or Boost.Regex
struct RegexpExtractEngine {
    std::unique_ptr<re2::RE2> re2_regex;
    std::unique_ptr<boost::regex> boost_regex;

    bool is_boost() const { return boost_regex != nullptr; }
    bool is_re2() const { return re2_regex != nullptr; }

    // Try to compile with RE2 first, fallback to Boost.Regex if RE2 fails
    static bool compile(const StringRef& pattern, std::string* error_str,
                        RegexpExtractEngine& engine, bool enable_extended_regex) {
        re2::RE2::Options options;
        options.set_log_errors(false); // avoid RE2 printing to stderr; we handle errors ourselves
        options.set_dot_nl(true); // make '.' match '\n' by default, consistent with REGEXP/LIKE
        engine.re2_regex =
                std::make_unique<re2::RE2>(re2::StringPiece(pattern.data, pattern.size), options);

        if (engine.re2_regex->ok()) {
            return true;
        } else if (!enable_extended_regex) {
            *error_str = fmt::format(
                    "Invalid regex pattern: {}. Error: {}. If you need advanced regex features, "
                    "try setting enable_extended_regex=true",
                    std::string(pattern.data, pattern.size), engine.re2_regex->error());
            return false;
        }

        // RE2 failed, try Boost.Regex for advanced features like zero-width assertions
        engine.re2_regex.reset();
        try {
            boost::regex::flag_type flags = boost::regex::normal;
            engine.boost_regex = std::make_unique<boost::regex>(pattern.data,
                                                                pattern.data + pattern.size, flags);
            return true;
        } catch (const boost::regex_error& e) {
            if (error_str) {
                *error_str = fmt::format("Invalid regex pattern: {}. Error: {}",
                                         std::string(pattern.data, pattern.size), e.what());
            }
            return false;
        }
    }

    // Get number of capturing groups
    int number_of_capturing_groups() const {
        if (is_re2()) {
            return re2_regex->NumberOfCapturingGroups();
        } else if (is_boost()) {
            return static_cast<int>(boost_regex->mark_count());
        }
        return 0;
    }

    // Match function for extraction
    bool match_and_extract(const char* data, size_t size, int index, std::string& result) const {
        if (is_re2()) {
            int max_matches = 1 + re2_regex->NumberOfCapturingGroups();
            if (index >= max_matches) {
                return false;
            }
            std::vector<re2::StringPiece> matches(max_matches);
            bool success = re2_regex->Match(re2::StringPiece(data, size), 0, size,
                                            re2::RE2::UNANCHORED, matches.data(), max_matches);
            if (success && index < matches.size()) {
                const re2::StringPiece& match = matches[index];
                result.assign(match.data(), match.size());
                return true;
            }
            return false;
        } else if (is_boost()) {
            boost::cmatch matches;
            bool success = boost::regex_search(data, data + size, matches, *boost_regex);
            if (success && index < matches.size()) {
                result = matches[index].str();
                return true;
            }
            return false;
        }
        return false;
    }

    // Match all occurrences and extract the first capturing group
    void match_all_and_extract(const char* data, size_t size,
                               std::vector<std::string>& results) const {
        if (is_re2()) {
            int max_matches = 1 + re2_regex->NumberOfCapturingGroups();
            if (max_matches < 2) {
                return; // No capturing groups
            }

            size_t pos = 0;
            while (pos < size) {
                const char* str_pos = data + pos;
                size_t str_size = size - pos;
                std::vector<re2::StringPiece> matches(max_matches);
                bool success = re2_regex->Match(re2::StringPiece(str_pos, str_size), 0, str_size,
                                                re2::RE2::UNANCHORED, matches.data(), max_matches);
                if (!success) {
                    break;
                }
                if (matches[0].empty()) {
                    pos += 1;
                    continue;
                }
                // Extract first capturing group
                if (matches.size() > 1 && !matches[1].empty()) {
                    results.emplace_back(matches[1].data(), matches[1].size());
                }
                // Move position forward
                auto offset = std::string(str_pos, str_size)
                                      .find(std::string(matches[0].data(), matches[0].size()));
                pos += offset + matches[0].size();
            }
        } else if (is_boost()) {
            const char* search_start = data;
            const char* search_end = data + size;
            boost::match_results<const char*> matches;

            while (boost::regex_search(search_start, search_end, matches, *boost_regex)) {
                if (matches.size() > 1 && matches[1].matched) {
                    results.emplace_back(matches[1].str());
                }
                if (matches[0].length() == 0) {
                    if (search_start == search_end) {
                        break;
                    }
                    search_start += 1;
                } else {
                    search_start = matches[0].second;
                }
            }
        }
    }
};

struct RegexpCountImpl {
    using StringColumnView = ColumnView<TYPE_STRING>;

    static void execute_impl(FunctionContext* context, ColumnPtr argument_columns[],
                             size_t input_rows_count, ColumnInt32::Container& result_data) {
        auto str_col = StringColumnView::create(argument_columns[0]);
        auto pattern_col = StringColumnView::create(argument_columns[1]);
        for (size_t i = 0; i < input_rows_count; ++i) {
            DCHECK(!str_col.is_null_at(i));
            DCHECK(!pattern_col.is_null_at(i));
            result_data[i] = _execute_inner_loop(context, str_col, pattern_col, i);
        }
    }
    static int _execute_inner_loop(FunctionContext* context, const StringColumnView& str_col,
                                   const StringColumnView& pattern_col, const size_t index_now) {
        re2::RE2* re = reinterpret_cast<re2::RE2*>(
                context->get_function_state(FunctionContext::THREAD_LOCAL));
        std::unique_ptr<re2::RE2> scoped_re;
        if (re == nullptr) {
            std::string error_str;
            const auto pattern = pattern_col.value_at(index_now);
            bool st = StringFunctions::compile_regex(pattern, &error_str, StringRef(), StringRef(),
                                                     scoped_re);
            if (!st) {
                context->add_warning(error_str.c_str());
                throw Exception(Status::InvalidArgument(error_str));
                return 0;
            }
            re = scoped_re.get();
        }

        const auto str = str_col.value_at(index_now);
        int count = 0;
        size_t pos = 0;
        while (pos < str.size) {
            auto str_pos = str.data + pos;
            auto str_size = str.size - pos;
            re2::StringPiece str_sp_current = re2::StringPiece(str_pos, str_size);
            re2::StringPiece match;

            bool success = re->Match(str_sp_current, 0, str_size, re2::RE2::UNANCHORED, &match, 1);
            if (!success) {
                break;
            }
            if (match.empty()) {
                pos += 1;
                continue;
            }
            count++;
            size_t match_start = match.data() - str_sp_current.data();
            pos += match_start + match.size();
        }

        return count;
    }
};

class FunctionRegexpCount : public IFunction {
public:
    static constexpr auto name = "regexp_count";

    static FunctionPtr create() { return std::make_shared<FunctionRegexpCount>(); }

    String get_name() const override { return name; }

    size_t get_number_of_arguments() const override { return 2; }

    DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {
        return std::make_shared<DataTypeInt32>();
    }

    Status open(FunctionContext* context, FunctionContext::FunctionStateScope scope) override {
        if (scope == FunctionContext::THREAD_LOCAL) {
            if (context->is_col_constant(1)) {
                DCHECK(!context->get_function_state(scope));
                const auto pattern_col = context->get_constant_col(1)->column_ptr;
                const auto& pattern = pattern_col->get_data_at(0);
                if (pattern.size == 0) {
                    return Status::OK();
                }

                std::string error_str;
                std::unique_ptr<re2::RE2> scoped_re;
                bool st = StringFunctions::compile_regex(pattern, &error_str, StringRef(),
                                                         StringRef(), scoped_re);
                if (!st) {
                    context->set_error(error_str.c_str());
                    return Status::InvalidArgument(error_str);
                }
                std::shared_ptr<re2::RE2> re(scoped_re.release());
                context->set_function_state(scope, re);
            }
        }
        return Status::OK();
    }

    Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments,
                        uint32_t result, size_t input_rows_count) const override {
        auto result_data_column = ColumnInt32::create(input_rows_count);
        auto& result_data = result_data_column->get_data();

        ColumnPtr argument_columns[2];

        argument_columns[0] = block.get_by_position(arguments[0]).column;
        argument_columns[1] = block.get_by_position(arguments[1]).column;
        RegexpCountImpl::execute_impl(context, argument_columns, input_rows_count, result_data);

        block.get_by_position(result).column = std::move(result_data_column);
        return Status::OK();
    }
};

struct ThreeParamTypes {
    static DataTypes get_variadic_argument_types() {
        return {std::make_shared<DataTypeString>(), std::make_shared<DataTypeString>(),
                std::make_shared<DataTypeString>()};
    }
};

struct FourParamTypes {
    static DataTypes get_variadic_argument_types() {
        return {std::make_shared<DataTypeString>(), std::make_shared<DataTypeString>(),
                std::make_shared<DataTypeString>(), std::make_shared<DataTypeString>()};
    }
};

// template FunctionRegexpFunctionality is used for regexp_replace/regexp_replace_one
template <typename Impl, typename ParamTypes>
class FunctionRegexpReplace : public IFunction {
public:
    static constexpr auto name = Impl::name;

    static FunctionPtr create() { return std::make_shared<FunctionRegexpReplace>(); }

    String get_name() const override { return name; }

    size_t get_number_of_arguments() const override {
        return get_variadic_argument_types_impl().size();
    }

    bool is_variadic() const override { return true; }

    DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {
        return make_nullable(std::make_shared<DataTypeString>());
    }

    DataTypes get_variadic_argument_types_impl() const override {
        return ParamTypes::get_variadic_argument_types();
    }

    Status open(FunctionContext* context, FunctionContext::FunctionStateScope scope) override {
        if (scope == FunctionContext::THREAD_LOCAL) {
            if (context->is_col_constant(1)) {
                DCHECK(!context->get_function_state(scope));
                const auto pattern_col = context->get_constant_col(1)->column_ptr;
                const auto& pattern = pattern_col->get_data_at(0);
                if (pattern.size == 0) {
                    return Status::OK();
                }

                std::string error_str;
                std::unique_ptr<re2::RE2> scoped_re;
                StringRef options_value;
                if constexpr (std::is_same_v<FourParamTypes, ParamTypes>) {
                    DCHECK_EQ(context->get_num_args(), 4);
                    DCHECK(context->is_col_constant(3));
                    const auto options_col = context->get_constant_col(3)->column_ptr;
                    options_value = options_col->get_data_at(0);
                }

                bool st = StringFunctions::compile_regex(pattern, &error_str, StringRef(),
                                                         options_value, scoped_re);
                if (!st) {
                    context->set_error(error_str.c_str());
                    return Status::InvalidArgument(error_str);
                }
                std::shared_ptr<re2::RE2> re(scoped_re.release());
                context->set_function_state(scope, re);
            }
        }
        return Status::OK();
    }

    Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments,
                        uint32_t result, size_t input_rows_count) const override {
        size_t argument_size = arguments.size();

        auto result_null_map = ColumnUInt8::create(input_rows_count, 0);
        auto result_data_column = ColumnString::create();
        auto& result_data = result_data_column->get_chars();
        auto& result_offset = result_data_column->get_offsets();
        result_offset.resize(input_rows_count);

        bool col_const[3];
        ColumnPtr argument_columns[3];
        for (int i = 0; i < 3; ++i) {
            col_const[i] = is_column_const(*block.get_by_position(arguments[i]).column);
        }
        argument_columns[0] = col_const[0] ? static_cast<const ColumnConst&>(
                                                     *block.get_by_position(arguments[0]).column)
                                                     .convert_to_full_column()
                                           : block.get_by_position(arguments[0]).column;

        default_preprocess_parameter_columns(argument_columns, col_const, {1, 2}, block, arguments);

        StringRef options_value;
        if (col_const[1] && col_const[2]) {
            Impl::execute_impl_const_args(context, argument_columns, options_value,
                                          input_rows_count, result_data, result_offset,
                                          result_null_map->get_data());
        } else {
            // the options have check in FE, so is always const, and get idx of 0
            if (argument_size == 4) {
                options_value = block.get_by_position(arguments[3]).column->get_data_at(0);
            }
            Impl::execute_impl(context, argument_columns, options_value, input_rows_count,
                               result_data, result_offset, result_null_map->get_data());
        }

        block.get_by_position(result).column =
                ColumnNullable::create(std::move(result_data_column), std::move(result_null_map));
        return Status::OK();
    }
};

template <bool ReplaceOne>
struct RegexpReplaceImpl {
    static constexpr auto name = ReplaceOne ? "regexp_replace_one" : "regexp_replace";

    static void execute_impl(FunctionContext* context, ColumnPtr argument_columns[],
                             const StringRef& options_value, size_t input_rows_count,
                             ColumnString::Chars& result_data, ColumnString::Offsets& result_offset,
                             NullMap& null_map) {
        const auto* str_col = check_and_get_column<ColumnString>(argument_columns[0].get());
        const auto* pattern_col = check_and_get_column<ColumnString>(argument_columns[1].get());
        const auto* replace_col = check_and_get_column<ColumnString>(argument_columns[2].get());

        for (size_t i = 0; i < input_rows_count; ++i) {
            _execute_inner_loop<false>(context, str_col, pattern_col, replace_col, options_value,
                                       result_data, result_offset, null_map, i);
        }
    }

    static void execute_impl_const_args(FunctionContext* context, ColumnPtr argument_columns[],
                                        const StringRef& options_value, size_t input_rows_count,
                                        ColumnString::Chars& result_data,
                                        ColumnString::Offsets& result_offset, NullMap& null_map) {
        const auto* str_col = check_and_get_column<ColumnString>(argument_columns[0].get());
        const auto* pattern_col = check_and_get_column<ColumnString>(argument_columns[1].get());
        const auto* replace_col = check_and_get_column<ColumnString>(argument_columns[2].get());

        for (size_t i = 0; i < input_rows_count; ++i) {
            _execute_inner_loop<true>(context, str_col, pattern_col, replace_col, options_value,
                                      result_data, result_offset, null_map, i);
        }
    }

    template <bool Const>
    static void _execute_inner_loop(FunctionContext* context, const ColumnString* str_col,
                                    const ColumnString* pattern_col,
                                    const ColumnString* replace_col, const StringRef& options_value,
                                    ColumnString::Chars& result_data,
                                    ColumnString::Offsets& result_offset, NullMap& null_map,
                                    const size_t index_now) {
        re2::RE2* re = reinterpret_cast<re2::RE2*>(
                context->get_function_state(FunctionContext::THREAD_LOCAL));
        std::unique_ptr<re2::RE2> scoped_re; // destroys re if state->re is nullptr
        if (re == nullptr) {
            std::string error_str;
            const auto& pattern = pattern_col->get_data_at(index_check_const(index_now, Const));
            bool st = StringFunctions::compile_regex(pattern, &error_str, StringRef(),
                                                     options_value, scoped_re);
            if (!st) {
                context->add_warning(error_str.c_str());
                StringOP::push_null_string(index_now, result_data, result_offset, null_map);
                return;
            }
            re = scoped_re.get();
        }

        re2::StringPiece replace_str = re2::StringPiece(
                replace_col->get_data_at(index_check_const(index_now, Const)).to_string_view());

        std::string result_str(str_col->get_data_at(index_now).to_string());
        if constexpr (ReplaceOne) {
            re2::RE2::Replace(&result_str, *re, replace_str);
        } else {
            re2::RE2::GlobalReplace(&result_str, *re, replace_str);
        }
        StringOP::push_value_string(result_str, index_now, result_data, result_offset);
    }
};

template <bool ReturnNull>
struct RegexpExtractImpl {
    static constexpr auto name = ReturnNull ? "regexp_extract_or_null" : "regexp_extract";
    static constexpr size_t num_args = 3;
    static constexpr size_t PATTERN_ARG_IDX = 1;

    static DataTypePtr return_type() { return make_nullable(std::make_shared<DataTypeString>()); }

    static Status execute(FunctionContext* context, Block& block, const ColumnNumbers& arguments,
                          uint32_t result, size_t input_rows_count) {
        bool col_const[3];
        ColumnPtr argument_columns[3];
        for (int i = 0; i < 3; ++i) {
            col_const[i] = is_column_const(*block.get_by_position(arguments[i]).column);
        }
        argument_columns[0] = col_const[0] ? static_cast<const ColumnConst&>(
                                                     *block.get_by_position(arguments[0]).column)
                                                     .convert_to_full_column()
                                           : block.get_by_position(arguments[0]).column;

        auto result_null_map = ColumnUInt8::create(input_rows_count, 0);
        auto result_data_column = ColumnString::create();
        auto& result_data = result_data_column->get_chars();
        auto& result_offset = result_data_column->get_offsets();
        result_offset.resize(input_rows_count);
        auto& null_map = result_null_map->get_data();

        default_preprocess_parameter_columns(argument_columns, col_const, {1, 2}, block, arguments);

        if (col_const[1] && col_const[2]) {
            _execute_loop<true>(context, argument_columns, input_rows_count, result_data,
                                result_offset, null_map);
        } else {
            _execute_loop<false>(context, argument_columns, input_rows_count, result_data,
                                 result_offset, null_map);
        }

        block.get_by_position(result).column =
                ColumnNullable::create(std::move(result_data_column), std::move(result_null_map));
        return Status::OK();
    }

private:
    template <bool Const>
    static void _execute_loop(FunctionContext* context, ColumnPtr argument_columns[],
                              size_t input_rows_count, ColumnString::Chars& result_data,
                              ColumnString::Offsets& result_offset, NullMap& null_map) {
        const auto* str_col = check_and_get_column<ColumnString>(argument_columns[0].get());
        const auto* pattern_col = check_and_get_column<ColumnString>(argument_columns[1].get());
        const auto* index_col = check_and_get_column<ColumnInt64>(argument_columns[2].get());

        if constexpr (Const) {
            const auto& index_data = index_col->get_int(0);
            if (index_data < 0) {
                for (size_t i = 0; i < input_rows_count; ++i) {
                    ReturnNull ? StringOP::push_null_string(i, result_data, result_offset, null_map)
                               : StringOP::push_empty_string(i, result_data, result_offset);
                }
                return;
            }
            for (size_t i = 0; i < input_rows_count; ++i) {
                _execute_inner_loop<true>(context, str_col, pattern_col, index_data, result_data,
                                          result_offset, null_map, i);
            }
        } else {
            for (size_t i = 0; i < input_rows_count; ++i) {
                const auto& index_data = index_col->get_int(i);
                if (index_data < 0) {
                    ReturnNull ? StringOP::push_null_string(i, result_data, result_offset, null_map)
                               : StringOP::push_empty_string(i, result_data, result_offset);
                    continue;
                }
                _execute_inner_loop<false>(context, str_col, pattern_col, index_data, result_data,
                                           result_offset, null_map, i);
            }
        }
    }

    template <bool Const>
    static void _execute_inner_loop(FunctionContext* context, const ColumnString* str_col,
                                    const ColumnString* pattern_col, const Int64 index_data,
                                    ColumnString::Chars& result_data,
                                    ColumnString::Offsets& result_offset, NullMap& null_map,
                                    const size_t index_now) {
        auto* engine = reinterpret_cast<RegexpExtractEngine*>(
                context->get_function_state(FunctionContext::THREAD_LOCAL));
        std::unique_ptr<RegexpExtractEngine> scoped_engine;

        if (engine == nullptr) {
            std::string error_str;
            const auto& pattern = pattern_col->get_data_at(index_check_const(index_now, Const));
            scoped_engine = std::make_unique<RegexpExtractEngine>();
            bool st = RegexpExtractEngine::compile(pattern, &error_str, *scoped_engine,
                                                   context->state()->enable_extended_regex());
            if (!st) {
                context->add_warning(error_str.c_str());
                StringOP::push_null_string(index_now, result_data, result_offset, null_map);
                return;
            }
            engine = scoped_engine.get();
        }

        const auto& str = str_col->get_data_at(index_now);

        int max_matches = 1 + engine->number_of_capturing_groups();
        if (index_data >= max_matches) {
            ReturnNull ? StringOP::push_null_string(index_now, result_data, result_offset, null_map)
                       : StringOP::push_empty_string(index_now, result_data, result_offset);
            return;
        }

        std::string match_result;
        bool success = engine->match_and_extract(str.data, str.size, static_cast<int>(index_data),
                                                 match_result);

        if (!success) {
            ReturnNull ? StringOP::push_null_string(index_now, result_data, result_offset, null_map)
                       : StringOP::push_empty_string(index_now, result_data, result_offset);
            return;
        }

        StringOP::push_value_string(std::string_view(match_result.data(), match_result.size()),
                                    index_now, result_data, result_offset);
    }
};

// Output handler for existing string-formatted result: "['a','b']"
struct RegexpExtractAllStringOutput {
    static constexpr const char* func_name = "regexp_extract_all";
    static DataTypePtr return_type() { return make_nullable(std::make_shared<DataTypeString>()); }

    ColumnString::Chars& result_data;
    ColumnString::Offsets& result_offset;

    void push_empty(size_t index) {
        StringOP::push_empty_string(index, result_data, result_offset);
    }
    void push_null(size_t index, NullMap& null_map) {
        StringOP::push_null_string(index, result_data, result_offset, null_map);
    }
    void push_matches(size_t index, const std::vector<std::string>& matches) {
        size_t total_size = 2; // '[' and ']'
        for (const auto& m : matches) {
            total_size += m.size() + 3; // "'xxx',"
        }

        size_t old_size = result_data.size();
        result_data.resize(old_size + total_size);
        char* pos = reinterpret_cast<char*>(&result_data[old_size]);

        *pos++ = '[';
        for (size_t j = 0; j < matches.size(); ++j) {
            if (j > 0) {
                *pos++ = ',';
            }
            *pos++ = '\'';
            memcpy(pos, matches[j].data(), matches[j].size());
            pos += matches[j].size();
            *pos++ = '\'';
        }
        *pos++ = ']';

        result_data.resize(old_size + static_cast<size_t>(pos - reinterpret_cast<char*>(
                                                                        &result_data[old_size])));
        result_offset[index] = static_cast<ColumnString::Offset>(result_data.size());
    }

    struct State {
        ColumnString::MutablePtr data_column;
        explicit State(size_t rows) : data_column(ColumnString::create()) {
            data_column->get_offsets().resize(rows);
        }
        RegexpExtractAllStringOutput create_handler() {
            return {.result_data = data_column->get_chars(),
                    .result_offset = data_column->get_offsets()};
        }
        ColumnPtr finalize(ColumnUInt8::MutablePtr null_map) {
            return ColumnNullable::create(std::move(data_column), std::move(null_map));
        }
    };
};

// Output handler for proper Array<Nullable<String>> result
struct RegexpExtractAllArrayOutput {
    static constexpr const char* func_name = "regexp_extract_all_array";
    static DataTypePtr return_type() {
        return make_nullable(
                std::make_shared<DataTypeArray>(make_nullable(std::make_shared<DataTypeString>())));
    }

    ColumnString& nested_col;
    ColumnArray::Offsets64& array_offsets;
    NullMap& nested_null_map;
    UInt64 current_offset = 0;

    void push_empty(size_t index) { array_offsets.push_back(current_offset); }
    void push_null(size_t index, NullMap& null_map) {
        null_map[index] = 1;
        array_offsets.push_back(current_offset);
    }
    void push_matches(size_t index, const std::vector<std::string>& matches) {
        for (const auto& m : matches) {
            nested_col.insert_data(m.data(), m.size());
            nested_null_map.push_back(0);
            current_offset++;
        }
        array_offsets.push_back(current_offset);
    }

    struct State {
        ColumnArray::MutablePtr array_column;
        ColumnNullable* nested_nullable;
        explicit State(size_t /*rows*/) {
            auto nullable_str = make_nullable(std::make_shared<DataTypeString>());
            array_column = ColumnArray::create(nullable_str->create_column(),
                                               ColumnArray::ColumnOffsets::create());
            nested_nullable = assert_cast<ColumnNullable*>(&array_column->get_data());
        }
        RegexpExtractAllArrayOutput create_handler() {
            return {.nested_col = assert_cast<ColumnString&>(nested_nullable->get_nested_column()),
                    .array_offsets = array_column->get_offsets(),
                    .nested_null_map = nested_nullable->get_null_map_data()};
        }
        ColumnPtr finalize(ColumnUInt8::MutablePtr null_map) {
            return ColumnNullable::create(std::move(array_column), std::move(null_map));
        }
    };
};

// Handler controls return type & column layout
template <typename Handler>
struct RegexpExtractAllImpl {
    static constexpr auto name = Handler::func_name;
    static constexpr size_t num_args = 2;
    static constexpr size_t PATTERN_ARG_IDX = 1;

    static DataTypePtr return_type() { return Handler::return_type(); }

    static Status execute(FunctionContext* context, Block& block, const ColumnNumbers& arguments,
                          uint32_t result, size_t input_rows_count) {
        bool col_const[2];
        ColumnPtr argument_columns[2];
        for (int i = 0; i < 2; ++i) {
            col_const[i] = is_column_const(*block.get_by_position(arguments[i]).column);
        }
        argument_columns[0] = col_const[0] ? static_cast<const ColumnConst&>(
                                                     *block.get_by_position(arguments[0]).column)
                                                     .convert_to_full_column()
                                           : block.get_by_position(arguments[0]).column;

        default_preprocess_parameter_columns(argument_columns, col_const, {1}, block, arguments);

        const auto* str_col = check_and_get_column<ColumnString>(argument_columns[0].get());
        const auto* pattern_col = check_and_get_column<ColumnString>(argument_columns[1].get());

        auto outer_null_map = ColumnUInt8::create(input_rows_count, 0);
        auto& null_map_data = outer_null_map->get_data();

        typename Handler::State state(input_rows_count);
        auto handler = state.create_handler();

        std::visit(
                [&](auto is_const) {
                    for (size_t i = 0; i < input_rows_count; ++i) {
                        if (null_map_data[i]) {
                            handler.push_null(i, null_map_data);
                            continue;
                        }
                        regexp_extract_all_inner_loop<is_const>(context, str_col, pattern_col,
                                                                handler, null_map_data, i);
                    }
                },
                make_bool_variant(col_const[1]));

        block.get_by_position(result).column = state.finalize(std::move(outer_null_map));
        return Status::OK();
    }

private:
    template <bool is_const>
    static void regexp_extract_all_inner_loop(FunctionContext* context, const ColumnString* str_col,
                                              const ColumnString* pattern_col, Handler& handler,
                                              NullMap& null_map, const size_t index_now) {
        auto* engine = reinterpret_cast<RegexpExtractEngine*>(
                context->get_function_state(FunctionContext::THREAD_LOCAL));
        std::unique_ptr<RegexpExtractEngine> scoped_engine;

        if (engine == nullptr) {
            std::string error_str;
            const auto& pattern = pattern_col->get_data_at(index_check_const(index_now, is_const));
            scoped_engine = std::make_unique<RegexpExtractEngine>();
            bool st = RegexpExtractEngine::compile(pattern, &error_str, *scoped_engine,
                                                   context->state()->enable_extended_regex());
            if (!st) {
                context->add_warning(error_str.c_str());
                handler.push_null(index_now, null_map);
                return;
            }
            engine = scoped_engine.get();
        }

        if (engine->number_of_capturing_groups() == 0) {
            handler.push_empty(index_now);
            return;
        }
        const auto& str = str_col->get_data_at(index_now);
        std::vector<std::string> res_matches;
        engine->match_all_and_extract(str.data, str.size, res_matches);

        if (res_matches.empty()) {
            handler.push_empty(index_now);
            return;
        }
        handler.push_matches(index_now, res_matches);
    }
};

// template FunctionRegexpFunctionality is used for regexp_xxxx series functions, not for regexp match.
template <typename Impl>
class FunctionRegexpFunctionality : public IFunction {
public:
    static constexpr auto name = Impl::name;

    static FunctionPtr create() { return std::make_shared<FunctionRegexpFunctionality>(); }

    String get_name() const override { return name; }

    size_t get_number_of_arguments() const override { return Impl::num_args; }

    DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {
        return Impl::return_type();
    }

    Status open(FunctionContext* context, FunctionContext::FunctionStateScope scope) override {
        if (scope == FunctionContext::THREAD_LOCAL) {
            if (context->is_col_constant(Impl::PATTERN_ARG_IDX)) {
                DCHECK(!context->get_function_state(scope));
                const auto pattern_col =
                        context->get_constant_col(Impl::PATTERN_ARG_IDX)->column_ptr;
                const auto& pattern = pattern_col->get_data_at(0);
                if (pattern.size == 0) {
                    return Status::OK();
                }

                std::string error_str;
                auto engine = std::make_shared<RegexpExtractEngine>();
                bool st = RegexpExtractEngine::compile(pattern, &error_str, *engine,
                                                       context->state()->enable_extended_regex());

                if (!st) {
                    context->set_error(error_str.c_str());
                    return Status::InvalidArgument(error_str);
                }
                context->set_function_state(scope, engine);
            }
        }
        return Status::OK();
    }

    Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments,
                        uint32_t result, size_t input_rows_count) const override {
        return Impl::execute(context, block, arguments, result, input_rows_count);
    }
};

void register_function_regexp_extract(SimpleFunctionFactory& factory) {
    factory.register_function<FunctionRegexpReplace<RegexpReplaceImpl<false>, ThreeParamTypes>>();
    factory.register_function<FunctionRegexpReplace<RegexpReplaceImpl<false>, FourParamTypes>>();
    factory.register_function<FunctionRegexpReplace<RegexpReplaceImpl<true>, ThreeParamTypes>>();
    factory.register_function<FunctionRegexpReplace<RegexpReplaceImpl<true>, FourParamTypes>>();
    factory.register_function<FunctionRegexpFunctionality<RegexpExtractImpl<true>>>();
    factory.register_function<FunctionRegexpFunctionality<RegexpExtractImpl<false>>>();
    factory.register_function<
            FunctionRegexpFunctionality<RegexpExtractAllImpl<RegexpExtractAllStringOutput>>>();
    factory.register_function<
            FunctionRegexpFunctionality<RegexpExtractAllImpl<RegexpExtractAllArrayOutput>>>();
    factory.register_function<FunctionRegexpCount>();
}

} // namespace doris

Coverage Report

Created: 2026-07-13 02:04