Skip to content

[Feature](func) Support the third argument for regexp_extract_all and regexp_extract_all_array #66050

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ccl125 wants to merge 1 commit into
base: master
Choose a base branch
from
Open

[Feature](func) Support the third argument for regexp_extract_all and regexp_extract_all_array #66050

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
54 changes: 35 additions & 19 deletions be/src/exprs/function/function_regexp.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,13 +136,17 @@ struct RegexpExtractEngine {
return false;
}

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

size_t pos = 0;
Expand All @@ -159,9 +163,9 @@ struct RegexpExtractEngine {
pos += 1;
continue;
}
// Extract first capturing group
if (matches.size() > 1 && !matches[1].empty()) {
results.emplace_back(matches[1].data(), matches[1].size());
// Extract the capturing group with the given index
if (static_cast<size_t>(index) < matches.size() && !matches[index].empty()) {
results.emplace_back(matches[index].data(), matches[index].size());
}
// Move position forward
auto offset = std::string(str_pos, str_size)
Expand All @@ -174,8 +178,8 @@ struct RegexpExtractEngine {
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 (static_cast<size_t>(index) < matches.size() && matches[index].matched) {
results.emplace_back(matches[index].str());
}
if (matches[0].length() == 0) {
if (search_start == search_end) {
Expand Down Expand Up @@ -716,27 +720,28 @@ struct RegexpExtractAllArrayOutput {
template <typename Handler>
struct RegexpExtractAllImpl {
static constexpr auto name = Handler::func_name;
static constexpr size_t num_args = 2;
static constexpr size_t num_args = 3;
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) {
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}, block, arguments);
default_preprocess_parameter_columns(argument_columns, col_const, {1, 2}, 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());
const auto* index_col = check_and_get_column<ColumnInt64>(argument_columns[2].get());

auto outer_null_map = ColumnUInt8::create(input_rows_count, 0);
auto& null_map_data = outer_null_map->get_data();
Expand All @@ -751,11 +756,13 @@ struct RegexpExtractAllImpl {
handler.push_null(i, null_map_data);
continue;
}
const auto index_data = index_col->get_int(index_check_const(i, is_const));
regexp_extract_all_inner_loop<is_const>(context, str_col, pattern_col,
handler, null_map_data, i);
index_data, handler, null_map_data,
i);
}
},
make_bool_variant(col_const[1]));
make_bool_variant(col_const[1] && col_const[2]));

block.get_by_position(result).column = state.finalize(std::move(outer_null_map));
return Status::OK();
Expand All @@ -764,8 +771,14 @@ struct RegexpExtractAllImpl {
private:
template <bool is_const>
static void regexp_extract_all_inner_loop(FunctionContext* context, const ColumnString* str_col,
const ColumnString* pattern_col, Handler& handler,
const ColumnString* pattern_col,
const Int64 index_data, Handler& handler,
NullMap& null_map, const size_t index_now) {
if (index_data < 0) {
handler.push_null(index_now, null_map);
return;
}

auto* engine = reinterpret_cast<RegexpExtractEngine*>(
context->get_function_state(FunctionContext::THREAD_LOCAL));
std::unique_ptr<RegexpExtractEngine> scoped_engine;
Expand All @@ -784,13 +797,16 @@ struct RegexpExtractAllImpl {
engine = scoped_engine.get();
}

if (engine->number_of_capturing_groups() == 0) {
// index 0 extracts the whole match, so patterns without capturing groups are
// only rejected when a positive group index is requested.
if (index_data > engine->number_of_capturing_groups()) {
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);
engine->match_all_and_extract(str.data, str.size, static_cast<int>(index_data),
res_matches);

if (res_matches.empty()) {
handler.push_empty(index_now);
Expand Down
117 changes: 85 additions & 32 deletions be/test/exprs/function/function_like_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -221,27 +221,53 @@ TEST(FunctionLikeTest, regexp_extract_all) {
std::string func_name = "regexp_extract_all";

DataSet data_set = {
{{std::string("x=a3&x=18abc&x=2&y=3&x=4&x=17bcd"), std::string("x=([0-9]+)([a-z]+)")},
{{std::string("x=a3&x=18abc&x=2&y=3&x=4&x=17bcd"), std::string("x=([0-9]+)([a-z]+)"),
(int64_t)1},
std::string("['18','17']")},
{{std::string("x=a3&x=18abc&x=2&y=3&x=4"), std::string("^x=([a-z]+)([0-9]+)")},
{{std::string("x=a3&x=18abc&x=2&y=3&x=4"), std::string("^x=([a-z]+)([0-9]+)"), (int64_t)1},
std::string("['a']")},
{{std::string("http://a.m.baidu.com/i41915173660.htm"), std::string("i([0-9]+)")},
{{std::string("http://a.m.baidu.com/i41915173660.htm"), std::string("i([0-9]+)"),
(int64_t)1},
std::string("['41915173660']")},
{{std::string("http://a.m.baidu.com/i41915i73660.htm"), std::string("i([0-9]+)")},
{{std::string("http://a.m.baidu.com/i41915i73660.htm"), std::string("i([0-9]+)"),
(int64_t)1},
std::string("['41915','73660']")},

{{std::string("hitdecisiondlist"), std::string("(i)(.*?)(e)")}, std::string("['i']")},
{{std::string("hitdecisioendlist"), std::string("(i)(.*?)(e)")},
{{std::string("hitdecisiondlist"), std::string("(i)(.*?)(e)"), (int64_t)1},
std::string("['i']")},
{{std::string("hitdecisioendlist"), std::string("(i)(.*?)(e)"), (int64_t)1},
std::string("['i','i']")},
{{std::string("hitdecisioendliset"), std::string("(i)(.*?)(e)")},
{{std::string("hitdecisioendliset"), std::string("(i)(.*?)(e)"), (int64_t)1},
std::string("['i','i','i']")},

// group index 0 extracts the whole match
{{std::string("x=a3&x=18abc&x=2&y=3&x=4&x=17bcd"), std::string("x=([0-9]+)([a-z]+)"),
(int64_t)0},
std::string("['x=18abc','x=17bcd']")},
{{std::string("hitdecisiondlist"), std::string("(i)(.*?)(e)"), (int64_t)0},
std::string("['itde']")},
// other capturing groups
{{std::string("x=a3&x=18abc&x=2&y=3&x=4&x=17bcd"), std::string("x=([0-9]+)([a-z]+)"),
(int64_t)2},
std::string("['abc','bcd']")},
{{std::string("hitdecisiondlist"), std::string("(i)(.*?)(e)"), (int64_t)2},
std::string("['td']")},
// group index out of range returns an empty result
{{std::string("hitdecisiondlist"), std::string("(i)(.*?)(e)"), (int64_t)3},
std::string("")},
// group index 0 also works for patterns without capturing groups
{{std::string("ab1cd22"), std::string("[0-9]+"), (int64_t)0},
std::string("['1','22']")},
// negative group index returns null
{{std::string("hitdecisiondlist"), std::string("(i)(.*?)(e)"), (int64_t)-1}, Null()},
// null
{{std::string("abc"), Null()}, Null()},
{{Null(), std::string("i([0-9]+)")}, Null()}};
{{std::string("abc"), Null(), (int64_t)1}, Null()},
{{Null(), std::string("i([0-9]+)"), (int64_t)1}, Null()}};

// pattern is constant value
InputTypeSet const_pattern_input_types = {PrimitiveType::TYPE_VARCHAR,
Consted {PrimitiveType::TYPE_VARCHAR}};
Consted {PrimitiveType::TYPE_VARCHAR},
PrimitiveType::TYPE_BIGINT};
for (const auto& line : data_set) {
DataSet const_pattern_dataset = {line};
static_cast<void>(check_function<DataTypeString, true>(func_name, const_pattern_input_types,
Expand All @@ -255,55 +281,70 @@ TEST(FunctionLikeTest, regexp_extract_all_array) {
auto return_type = make_nullable(
std::make_shared<DataTypeArray>(make_nullable(std::make_shared<DataTypeString>())));

auto run_case = [&](const std::string& str, const std::string& pattern,
auto bigint_type = std::make_shared<DataTypeInt64>();
auto run_case = [&](const std::string& str, const std::string& pattern, int64_t idx,
const std::string& expected, bool expect_null = false) {
auto col_str = ColumnString::create();
col_str->insert_data(str.data(), str.size());
auto col_pattern = ColumnString::create();
col_pattern->insert_data(pattern.data(), pattern.size());
auto col_idx = ColumnInt64::create();
col_idx->insert_value(idx);

Block block;
block.insert({std::move(col_str), str_type, "str"});
block.insert({ColumnConst::create(std::move(col_pattern), 1), str_type, "pattern"});
block.insert({ColumnConst::create(std::move(col_idx), 1), bigint_type, "idx"});
block.insert({nullptr, return_type, "result"});

ColumnsWithTypeAndName arg_cols = {block.get_by_position(0), block.get_by_position(1)};
ColumnsWithTypeAndName arg_cols = {block.get_by_position(0), block.get_by_position(1),
block.get_by_position(2)};
auto func =
SimpleFunctionFactory::instance().get_function(func_name, arg_cols, return_type);
ASSERT_TRUE(func != nullptr);

std::vector<DataTypePtr> arg_types = {str_type, str_type};
std::vector<DataTypePtr> arg_types = {str_type, str_type, bigint_type};
FunctionUtils fn_utils({}, arg_types, false);
auto* fn_ctx = fn_utils.get_fn_ctx();
fn_ctx->set_constant_cols(
{nullptr, std::make_shared<ColumnPtrWrapper>(block.get_by_position(1).column)});
{nullptr, std::make_shared<ColumnPtrWrapper>(block.get_by_position(1).column),
std::make_shared<ColumnPtrWrapper>(block.get_by_position(2).column)});

ASSERT_EQ(Status::OK(), func->open(fn_ctx, FunctionContext::FRAGMENT_LOCAL));
ASSERT_EQ(Status::OK(), func->open(fn_ctx, FunctionContext::THREAD_LOCAL));
ASSERT_EQ(Status::OK(), func->execute(fn_ctx, block, {0, 1}, 2, 1));
ASSERT_EQ(Status::OK(), func->execute(fn_ctx, block, {0, 1, 2}, 3, 1));

auto result_col = block.get_by_position(2).column;
auto result_col = block.get_by_position(3).column;
ASSERT_TRUE(result_col.get() != nullptr);
if (expect_null) {
EXPECT_TRUE(result_col->is_null_at(0));
} else {
ASSERT_FALSE(result_col->is_null_at(0));
auto result_str = return_type->to_string(*result_col, 0);
EXPECT_EQ(expected, result_str)
<< "input: '" << str << "', pattern: '" << pattern << "'";
<< "input: '" << str << "', pattern: '" << pattern << "', idx: " << idx;
}

static_cast<void>(func->close(fn_ctx, FunctionContext::THREAD_LOCAL));
static_cast<void>(func->close(fn_ctx, FunctionContext::FRAGMENT_LOCAL));
};

run_case("x=a3&x=18abc&x=2&y=3&x=4&x=17bcd", "x=([0-9]+)([a-z]+)", "[\"18\", \"17\"]");
run_case("x=a3&x=18abc&x=2&y=3&x=4", "^x=([a-z]+)([0-9]+)", "[\"a\"]");
run_case("http://a.m.baidu.com/i41915173660.htm", "i([0-9]+)", "[\"41915173660\"]");
run_case("http://a.m.baidu.com/i41915i73660.htm", "i([0-9]+)", "[\"41915\", \"73660\"]");
run_case("hitdecisiondlist", "(i)(.*?)(e)", "[\"i\"]");
run_case("no_match_here", "x=([0-9]+)", "[]");
run_case("abc", "([a-z]+)", "[\"abc\"]");
run_case("x=a3&x=18abc&x=2&y=3&x=4&x=17bcd", "x=([0-9]+)([a-z]+)", 1, "[\"18\", \"17\"]");
run_case("x=a3&x=18abc&x=2&y=3&x=4", "^x=([a-z]+)([0-9]+)", 1, "[\"a\"]");
run_case("http://a.m.baidu.com/i41915173660.htm", "i([0-9]+)", 1, "[\"41915173660\"]");
run_case("http://a.m.baidu.com/i41915i73660.htm", "i([0-9]+)", 1, "[\"41915\", \"73660\"]");
run_case("hitdecisiondlist", "(i)(.*?)(e)", 1, "[\"i\"]");
run_case("no_match_here", "x=([0-9]+)", 1, "[]");
run_case("abc", "([a-z]+)", 1, "[\"abc\"]");
// group index 0 extracts the whole match
run_case("x=a3&x=18abc&x=2&y=3&x=4&x=17bcd", "x=([0-9]+)([a-z]+)", 0,
"[\"x=18abc\", \"x=17bcd\"]");
// other capturing groups
run_case("x=a3&x=18abc&x=2&y=3&x=4&x=17bcd", "x=([0-9]+)([a-z]+)", 2, "[\"abc\", \"bcd\"]");
// group index out of range returns an empty array
run_case("hitdecisiondlist", "(i)(.*?)(e)", 3, "[]");
// negative group index returns null
run_case("hitdecisiondlist", "(i)(.*?)(e)", -1, "", true);

// Helper for testing null input propagation
auto nullable_str_type = make_nullable(str_type);
Expand Down Expand Up @@ -339,24 +380,31 @@ TEST(FunctionLikeTest, regexp_extract_all_array) {
Block block;
block.insert({col_str, str_col_type, "str"});
block.insert({col_pattern, pattern_col_type, "pattern"});
auto col_idx = ColumnInt64::create();
col_idx->insert_value(1);
block.insert({ColumnConst::create(std::move(col_idx), 1),
std::make_shared<DataTypeInt64>(), "idx"});
block.insert({nullptr, return_type, "result"});

ColumnsWithTypeAndName arg_cols = {block.get_by_position(0), block.get_by_position(1)};
ColumnsWithTypeAndName arg_cols = {block.get_by_position(0), block.get_by_position(1),
block.get_by_position(2)};
auto func =
SimpleFunctionFactory::instance().get_function(func_name, arg_cols, return_type);
ASSERT_TRUE(func != nullptr);

std::vector<DataTypePtr> arg_types = {str_col_type, pattern_col_type};
std::vector<DataTypePtr> arg_types = {str_col_type, pattern_col_type,
std::make_shared<DataTypeInt64>()};
FunctionUtils fn_utils({}, arg_types, false);
auto* fn_ctx = fn_utils.get_fn_ctx();
fn_ctx->set_constant_cols(
{nullptr, std::make_shared<ColumnPtrWrapper>(block.get_by_position(1).column)});
{nullptr, std::make_shared<ColumnPtrWrapper>(block.get_by_position(1).column),
std::make_shared<ColumnPtrWrapper>(block.get_by_position(2).column)});

ASSERT_EQ(Status::OK(), func->open(fn_ctx, FunctionContext::FRAGMENT_LOCAL));
ASSERT_EQ(Status::OK(), func->open(fn_ctx, FunctionContext::THREAD_LOCAL));
ASSERT_EQ(Status::OK(), func->execute(fn_ctx, block, {0, 1}, 2, 1));
ASSERT_EQ(Status::OK(), func->execute(fn_ctx, block, {0, 1, 2}, 3, 1));

EXPECT_TRUE(block.get_by_position(2).column->is_null_at(0))
EXPECT_TRUE(block.get_by_position(3).column->is_null_at(0))
<< "Expected null for null_str=" << null_str << " null_pattern=" << null_pattern;

static_cast<void>(func->close(fn_ctx, FunctionContext::THREAD_LOCAL));
Expand All @@ -374,21 +422,26 @@ TEST(FunctionLikeTest, regexp_extract_all_array) {
col_str->insert_data("abc", 3);
auto col_pattern = ColumnString::create();
col_pattern->insert_data("(", 1);
auto col_idx = ColumnInt64::create();
col_idx->insert_value(1);
Block block;
block.insert({std::move(col_str), str_type, "str"});
block.insert({ColumnConst::create(std::move(col_pattern), 1), str_type, "pattern"});
block.insert({ColumnConst::create(std::move(col_idx), 1), bigint_type, "idx"});
block.insert({nullptr, return_type, "result"});

ColumnsWithTypeAndName arg_cols = {block.get_by_position(0), block.get_by_position(1)};
ColumnsWithTypeAndName arg_cols = {block.get_by_position(0), block.get_by_position(1),
block.get_by_position(2)};
auto func =
SimpleFunctionFactory::instance().get_function(func_name, arg_cols, return_type);
ASSERT_TRUE(func != nullptr);

std::vector<DataTypePtr> arg_types = {str_type, str_type};
std::vector<DataTypePtr> arg_types = {str_type, str_type, bigint_type};
FunctionUtils fn_utils({}, arg_types, false);
auto* fn_ctx = fn_utils.get_fn_ctx();
fn_ctx->set_constant_cols(
{nullptr, std::make_shared<ColumnPtrWrapper>(block.get_by_position(1).column)});
{nullptr, std::make_shared<ColumnPtrWrapper>(block.get_by_position(1).column),
std::make_shared<ColumnPtrWrapper>(block.get_by_position(2).column)});

ASSERT_EQ(Status::OK(), func->open(fn_ctx, FunctionContext::FRAGMENT_LOCAL));
// Invalid pattern should cause open() to fail for THREAD_LOCAL scope
Expand Down
Loading
Loading