Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit a03a728

Browse files
zhangstar333feiniaofeiafei
zhangstar333
authored and
feiniaofeiafei
committed
[Feature](function) support split_by_regexp function (apache#38259)
## Proposed changes docs link: apache/doris-website#904 <!--Describe your changes.-->
1 parent 316072a commit a03a728

8 files changed

Lines changed: 606 additions & 0 deletions

File tree

‎be/src/vec/functions/function_split_by_regexp.cpp‎

Lines changed: 378 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,378 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
#include <fmt/format.h>
19+
#include <glog/logging.h>
20+
21+
#include "common/status.h"
22+
#include "vec/columns/column_array.h"
23+
#include "vec/columns/column_const.h"
24+
#include "vec/core/types.h"
25+
#include "vec/data_types/data_type_array.h"
26+
#include "vec/data_types/data_type_number.h"
27+
#include "vec/data_types/data_type_string.h"
28+
#include "vec/functions/function.h"
29+
#include "vec/functions/function_string.h"
30+
#include "vec/functions/simple_function_factory.h"
31+
32+
namespace doris::vectorized {
33+
34+
struct Match {
35+
std::string::size_type offset;
36+
std::string::size_type length;
37+
};
38+
39+
class RegexpSplit {
40+
public:
41+
void init(re2::RE2* re2, int32_t max_splits);
42+
void set(const char* pos, const char* end);
43+
bool get(const char*& token_begin, const char*& token_end);
44+
45+
private:
46+
const char* _pos;
47+
const char* _end;
48+
49+
std::int32_t _max_splits = 0;
50+
std::vector<Match> _matches;
51+
int32_t _splits;
52+
re2::RE2* _re2 = nullptr;
53+
unsigned _number_of_subpatterns = 0;
54+
55+
unsigned match(const char* subject, size_t subject_size, std::vector<Match>& matches,
56+
unsigned limit) const;
57+
};
58+
59+
unsigned RegexpSplit::match(const char* subject, size_t subject_size, std::vector<Match>& matches,
60+
unsigned limit) const {
61+
matches.clear();
62+
63+
if (limit == 0) {
64+
return 0;
65+
}
66+
67+
limit = std::min(limit, _number_of_subpatterns + 1);
68+
std::vector<re2::StringPiece> pieces(limit);
69+
70+
if (!_re2->Match({subject, subject_size}, 0, subject_size, re2::RE2::UNANCHORED, pieces.data(),
71+
limit)) {
72+
return 0;
73+
} else {
74+
matches.resize(limit);
75+
for (size_t i = 0; i < limit; ++i) {
76+
if (pieces[i].empty()) {
77+
matches[i].offset = std::string::npos;
78+
matches[i].length = 0;
79+
} else {
80+
matches[i].offset = pieces[i].data() - subject;
81+
matches[i].length = pieces[i].length();
82+
}
83+
}
84+
return limit;
85+
}
86+
}
87+
88+
void RegexpSplit::init(re2::RE2* re2, int32_t max_splits) {
89+
_max_splits = max_splits;
90+
_re2 = re2;
91+
if (_re2) {
92+
_number_of_subpatterns = _re2->NumberOfCapturingGroups();
93+
}
94+
}
95+
96+
// Called for each next string.
97+
void RegexpSplit::set(const char* pos, const char* end) {
98+
_pos = pos;
99+
_end = end;
100+
_splits = 0;
101+
}
102+
103+
// Get the next token, if any, or return false.
104+
bool RegexpSplit::get(const char*& token_begin, const char*& token_end) {
105+
if (!_re2) {
106+
if (_pos == _end) {
107+
return false;
108+
}
109+
110+
token_begin = _pos;
111+
if (_max_splits != -1) {
112+
if (_splits == _max_splits - 1) {
113+
token_end = _end;
114+
_pos = _end;
115+
return true;
116+
}
117+
}
118+
119+
_pos += 1;
120+
token_end = _pos;
121+
++_splits;
122+
} else {
123+
if (!_pos || _pos > _end) {
124+
return false;
125+
}
126+
127+
token_begin = _pos;
128+
if (_max_splits != -1) {
129+
if (_splits == _max_splits - 1) {
130+
token_end = _end;
131+
_pos = nullptr;
132+
return true;
133+
}
134+
}
135+
136+
if (!match(_pos, _end - _pos, _matches, _number_of_subpatterns + 1) ||
137+
!_matches[0].length) {
138+
token_end = _end;
139+
_pos = _end + 1;
140+
} else {
141+
token_end = _pos + _matches[0].offset;
142+
_pos = token_end + _matches[0].length;
143+
++_splits;
144+
}
145+
}
146+
147+
return true;
148+
}
149+
150+
template <typename Impl>
151+
class SplitByRegexp : public IFunction {
152+
public:
153+
static constexpr auto name = "split_by_regexp";
154+
155+
static FunctionPtr create() { return std::make_shared<SplitByRegexp>(); }
156+
157+
String get_name() const override { return name; }
158+
159+
size_t get_number_of_arguments() const override {
160+
return get_variadic_argument_types_impl().size();
161+
}
162+
163+
bool is_variadic() const override { return true; }
164+
165+
DataTypes get_variadic_argument_types_impl() const override {
166+
return Impl::get_variadic_argument_types();
167+
}
168+
169+
DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {
170+
DCHECK(is_string(arguments[0]))
171+
<< "first argument for function: " << name << " should be string"
172+
<< " and arguments[0] is " << arguments[0]->get_name();
173+
DCHECK(is_string(arguments[1]))
174+
<< "second argument for function: " << name << " should be string"
175+
<< " and arguments[1] is " << arguments[1]->get_name();
176+
auto nullable_string_type = make_nullable(std::make_shared<DataTypeString>());
177+
return std::make_shared<DataTypeArray>(nullable_string_type);
178+
}
179+
180+
Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments,
181+
size_t result, size_t input_rows_count) const override {
182+
return Impl::execute_impl(context, block, arguments, result, input_rows_count);
183+
}
184+
};
185+
186+
struct ExecuteImpl {
187+
using NullMapType = PaddedPODArray<UInt8>;
188+
static Status execute_impl(FunctionContext* context, Block& block,
189+
const ColumnNumbers& arguments, size_t result,
190+
size_t input_rows_count) {
191+
const auto& [first_column, left_const] =
192+
unpack_if_const(block.get_by_position(arguments[0]).column);
193+
const auto& [second_column, right_const] =
194+
unpack_if_const(block.get_by_position(arguments[1]).column);
195+
const auto& [three_column, three_is_const] =
196+
unpack_if_const(block.get_by_position(arguments[2]).column);
197+
auto limit_value = assert_cast<const ColumnInt32&>(*three_column).get_int(0);
198+
const auto& src_column = assert_cast<const ColumnString&>(*first_column);
199+
const auto& pattern_column = assert_cast<const ColumnString&>(*second_column);
200+
201+
auto nullable_string_type = make_nullable(std::make_shared<DataTypeString>());
202+
auto dest_column_ptr = ColumnArray::create(nullable_string_type->create_column(),
203+
ColumnArray::ColumnOffsets::create());
204+
IColumn* dest_nested_column = &dest_column_ptr->get_data();
205+
auto& dest_offsets = dest_column_ptr->get_offsets();
206+
DCHECK(dest_nested_column != nullptr);
207+
208+
NullMapType* dest_nested_null_map = nullptr;
209+
auto* dest_nullable_col = assert_cast<ColumnNullable*>(dest_nested_column);
210+
auto& dest_column_string =
211+
assert_cast<ColumnString&>(*(dest_nullable_col->get_nested_column_ptr()));
212+
dest_nested_null_map = &dest_nullable_col->get_null_map_column().get_data();
213+
RE2::Options opts;
214+
opts.set_never_nl(false);
215+
opts.set_dot_nl(true);
216+
// split_by_regexp(ColumnString, "xxx")
217+
if (right_const) {
218+
RETURN_IF_ERROR(_execute_constant_pattern(
219+
src_column, pattern_column.get_data_at(0), dest_column_string, dest_offsets,
220+
dest_nested_null_map, limit_value, input_rows_count, &opts));
221+
} else if (left_const) {
222+
// split_by_regexp("xxx", ColumnString)
223+
_execute_constant_src_string(src_column.get_data_at(0), pattern_column,
224+
dest_column_string, dest_offsets, dest_nested_null_map,
225+
limit_value, input_rows_count, &opts);
226+
} else {
227+
// split_by_regexp(ColumnString, ColumnString)
228+
_execute_vector_vector(src_column, pattern_column, dest_column_string, dest_offsets,
229+
dest_nested_null_map, limit_value, input_rows_count, &opts);
230+
}
231+
232+
block.replace_by_position(result, std::move(dest_column_ptr));
233+
return Status::OK();
234+
}
235+
236+
private:
237+
static Status _execute_constant_pattern(const ColumnString& src_column_string,
238+
const StringRef& pattern_ref,
239+
ColumnString& dest_column_string,
240+
ColumnArray::Offsets64& dest_offsets,
241+
NullMapType* dest_nested_null_map, Int64 limit_value,
242+
size_t input_rows_count, RE2::Options* opts) {
243+
const char* token_begin = nullptr;
244+
const char* token_end = nullptr;
245+
UInt64 index = 0;
246+
std::unique_ptr<re2::RE2> re2_ptr = nullptr;
247+
if (pattern_ref.size) {
248+
re2_ptr = std::make_unique<re2::RE2>(pattern_ref.to_string_view(), *opts);
249+
}
250+
if (!re2_ptr->ok()) {
251+
return Status::RuntimeError("Invalid pattern: {}", pattern_ref.debug_string());
252+
}
253+
RegexpSplit RegexpSplit;
254+
RegexpSplit.init(re2_ptr.get(), limit_value);
255+
for (int row = 0; row < input_rows_count; ++row) {
256+
auto str_data = src_column_string.get_data_at(row);
257+
RegexpSplit.set(str_data.begin(), str_data.end());
258+
while (RegexpSplit.get(token_begin, token_end)) {
259+
size_t token_size = token_end - token_begin;
260+
dest_column_string.insert_data(token_begin, token_size);
261+
dest_nested_null_map->push_back(false);
262+
index += 1;
263+
}
264+
dest_offsets.push_back(index);
265+
}
266+
return Status::OK();
267+
}
268+
269+
static void _execute_constant_src_string(const StringRef& str_ref,
270+
const ColumnString& pattern_column,
271+
ColumnString& dest_column_string,
272+
ColumnArray::Offsets64& dest_offsets,
273+
NullMapType* dest_nested_null_map, Int64 limit_value,
274+
size_t input_rows_count, RE2::Options* opts) {
275+
const char* token_begin = nullptr;
276+
const char* token_end = nullptr;
277+
UInt64 index = 0;
278+
RegexpSplit RegexpSplit;
279+
280+
for (int row = 0; row < input_rows_count; ++row) {
281+
std::unique_ptr<re2::RE2> re2_ptr = nullptr;
282+
auto pattern = pattern_column.get_data_at(row);
283+
if (pattern.size) {
284+
re2_ptr = std::make_unique<re2::RE2>(pattern.to_string_view(), *opts);
285+
if (!re2_ptr->ok()) {
286+
dest_column_string.insert_default();
287+
dest_nested_null_map->push_back(true);
288+
index += 1;
289+
dest_offsets.push_back(index);
290+
continue;
291+
}
292+
}
293+
294+
RegexpSplit.init(re2_ptr.get(), limit_value);
295+
RegexpSplit.set(str_ref.begin(), str_ref.end());
296+
while (RegexpSplit.get(token_begin, token_end)) {
297+
size_t token_size = token_end - token_begin;
298+
dest_column_string.insert_data(token_begin, token_size);
299+
dest_nested_null_map->push_back(false);
300+
index += 1;
301+
}
302+
dest_offsets.push_back(index);
303+
}
304+
}
305+
306+
static void _execute_vector_vector(const ColumnString& src_column_string,
307+
const ColumnString& pattern_column,
308+
ColumnString& dest_column_string,
309+
ColumnArray::Offsets64& dest_offsets,
310+
NullMapType* dest_nested_null_map, Int64 limit_value,
311+
size_t input_rows_count, RE2::Options* opts) {
312+
const char* token_begin = nullptr;
313+
const char* token_end = nullptr;
314+
UInt64 index = 0;
315+
RegexpSplit RegexpSplit;
316+
317+
for (int row = 0; row < input_rows_count; ++row) {
318+
std::unique_ptr<re2::RE2> re2_ptr = nullptr;
319+
auto str_data = src_column_string.get_data_at(row);
320+
auto pattern = pattern_column.get_data_at(row);
321+
if (pattern.size) {
322+
re2_ptr = std::make_unique<re2::RE2>(pattern.to_string_view(), *opts);
323+
if (!re2_ptr->ok()) {
324+
dest_column_string.insert_default();
325+
dest_nested_null_map->push_back(true);
326+
index += 1;
327+
dest_offsets.push_back(index);
328+
continue;
329+
}
330+
}
331+
RegexpSplit.init(re2_ptr.get(), limit_value);
332+
RegexpSplit.set(str_data.begin(), str_data.end());
333+
while (RegexpSplit.get(token_begin, token_end)) {
334+
size_t token_size = token_end - token_begin;
335+
dest_column_string.insert_data(token_begin, token_size);
336+
dest_nested_null_map->push_back(false);
337+
index += 1;
338+
}
339+
dest_offsets.push_back(index);
340+
}
341+
}
342+
};
343+
344+
struct TwoArgumentImpl {
345+
static DataTypes get_variadic_argument_types() {
346+
return {std::make_shared<DataTypeString>(), std::make_shared<DataTypeString>()};
347+
}
348+
349+
static Status execute_impl(FunctionContext* context, Block& block,
350+
const ColumnNumbers& arguments, size_t result,
351+
size_t input_rows_count) {
352+
DCHECK_EQ(arguments.size(), 2);
353+
auto max_limit = ColumnConst::create(ColumnInt32::create(1, -1), input_rows_count);
354+
block.insert({std::move(max_limit), std::make_shared<DataTypeInt32>(), "max_limit"});
355+
ColumnNumbers temp_arguments = {arguments[0], arguments[1], block.columns() - 1};
356+
return ExecuteImpl::execute_impl(context, block, temp_arguments, result, input_rows_count);
357+
}
358+
};
359+
360+
struct ThreeArgumentImpl {
361+
static DataTypes get_variadic_argument_types() {
362+
return {std::make_shared<DataTypeString>(), std::make_shared<DataTypeString>(),
363+
std::make_shared<DataTypeInt32>()};
364+
}
365+
static Status execute_impl(FunctionContext* context, Block& block,
366+
const ColumnNumbers& arguments, size_t result,
367+
size_t input_rows_count) {
368+
DCHECK_EQ(arguments.size(), 3);
369+
return ExecuteImpl::execute_impl(context, block, arguments, result, input_rows_count);
370+
}
371+
};
372+
373+
void register_function_split_by_regexp(SimpleFunctionFactory& factory) {
374+
factory.register_function<SplitByRegexp<TwoArgumentImpl>>();
375+
factory.register_function<SplitByRegexp<ThreeArgumentImpl>>();
376+
}
377+
378+
} // namespace doris::vectorized

‎be/src/vec/functions/simple_function_factory.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,7 @@ void register_function_tokenize(SimpleFunctionFactory& factory);
105105
void register_function_url(SimpleFunctionFactory& factory);
106106
void register_function_ip(SimpleFunctionFactory& factory);
107107
void register_function_multi_match(SimpleFunctionFactory& factory);
108+
void register_function_split_by_regexp(SimpleFunctionFactory& factory);
108109

109110
class SimpleFunctionFactory {
110111
using Creator = std::function<FunctionBuilderPtr()>;
@@ -292,6 +293,7 @@ class SimpleFunctionFactory {
292293
register_function_ignore(instance);
293294
register_function_variant_element(instance);
294295
register_function_multi_match(instance);
296+
register_function_split_by_regexp(instance);
295297
});
296298
return instance;
297299
}

0 commit comments

Comments
 (0)