database/doris
2
0
Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

[feature](function)add search functions: multi_search_all_positions & multi_match_any (#13763)

Browse Source 
Co-authored-by: yiliang qiu <yiliang.qiu@qq.com>
This commit is contained in:
Yiliang Qiu
2022-11-07 11:50:55 +08:00
committed by GitHub
parent 7ffe88b579
commit e8d2fb6778
14 changed files with 1143 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

236
be/src/vec/functions/functions_multi_string_position.cpp Normal file
View File

@ -0,0 +1,236 @@
// 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.
// This file is copied from
// https://github.com/ClickHouse/ClickHouse/blob/master/src/Functions/FunctionsMultiStringPosition.h
// and modified by Doris
#include "function.h"
#include "function_helpers.h"
#include "vec/columns/column_array.h"
#include "vec/columns/column_fixed_length_object.h"
#include "vec/columns/column_string.h"
#include "vec/columns/column_vector.h"
#include "vec/common/pod_array.h"
#include "vec/common/volnitsky.h"
#include "vec/data_types/data_type_array.h"
#include "vec/data_types/data_type_number.h"
#include "vec/data_types/data_type_string.h"
#include "vec/functions/simple_function_factory.h"
namespace doris::vectorized {
template <typename Impl>
class FunctionMultiStringPosition : public IFunction {
public:
static constexpr auto name = Impl::name;
static FunctionPtr create() { return std::make_shared<FunctionMultiStringPosition>(); }
String get_name() const override { return name; }
size_t get_number_of_arguments() const override { return 2; }
bool use_default_implementation_for_constants() const override { return true; }
DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {
return std::make_shared<DataTypeArray>(make_nullable(std::make_shared<DataTypeInt32>()));
}
Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments,
size_t result, size_t input_rows_count) override {
ColumnPtr haystack_ptr = block.get_by_position(arguments[0]).column;
ColumnPtr needles_ptr = block.get_by_position(arguments[1]).column;
const ColumnString* col_haystack_vector =
check_and_get_column<ColumnString>(&*haystack_ptr);
const ColumnConst* col_haystack_const =
check_and_get_column_const<ColumnString>(&*haystack_ptr);
const ColumnArray* col_needles_vector =
check_and_get_column<ColumnArray>(needles_ptr.get());
const ColumnConst* col_needles_const =
check_and_get_column_const<ColumnArray>(needles_ptr.get());
if (col_haystack_const && col_needles_vector)
return Status::InvalidArgument(
"function '{}' doesn't support search with non-constant needles "
"in constant haystack",
name);
using ResultType = typename Impl::ResultType;
auto col_res = ColumnVector<ResultType>::create();
auto col_offsets = ColumnArray::ColumnOffsets::create();
auto& vec_res = col_res->get_data();
auto& offsets_res = col_offsets->get_data();
Status status;
if (col_needles_const)
status = Impl::vector_constant(
col_haystack_vector->get_chars(), col_haystack_vector->get_offsets(),
col_needles_const->get_value<Array>(), vec_res, offsets_res);
else
status = Impl::vector_vector(col_haystack_vector->get_chars(),
col_haystack_vector->get_offsets(),
col_needles_vector->get_data(),
col_needles_vector->get_offsets(), vec_res, offsets_res);
if (!status.ok()) return status;
auto nullable_col =
ColumnNullable::create(std::move(col_res), ColumnUInt8::create(col_res->size(), 0));
block.get_by_position(result).column =
ColumnArray::create(std::move(nullable_col), std::move(col_offsets));
return status;
}
};
template <typename Impl>
struct FunctionMultiSearchAllPositionsImpl {
using ResultType = Int32;
static constexpr auto name = "multi_search_all_positions";
static Status vector_constant(const ColumnString::Chars& haystack_data,
const ColumnString::Offsets& haystack_offsets,
const Array& needles_arr, PaddedPODArray<Int32>& vec_res,
PaddedPODArray<UInt64>& offsets_res) {
if (needles_arr.size() > std::numeric_limits<UInt8>::max())
return Status::InvalidArgument(
"number of arguments for function {} doesn't match: "
"passed {}, should be at most 255",
name, needles_arr.size());
std::vector<StringRef> needles;
needles.reserve(needles_arr.size());
for (const auto& needle : needles_arr) needles.emplace_back(needle.get<StringRef>());
auto res_callback = [](const UInt8* start, const UInt8* end) -> Int32 {
return 1 + Impl::count_chars(reinterpret_cast<const char*>(start),
reinterpret_cast<const char*>(end));
};
auto searcher = Impl::create_multi_searcher(needles);
const size_t haystack_size = haystack_offsets.size();
const size_t needles_size = needles.size();
vec_res.resize(haystack_size * needles.size());
offsets_res.resize(haystack_size);
std::fill(vec_res.begin(), vec_res.end(), 0);
while (searcher.hasMoreToSearch()) {
size_t prev_haystack_offset = 0;
for (size_t j = 0, from = 0; j < haystack_size; ++j, from += needles_size) {
const auto* haystack = &haystack_data[prev_haystack_offset];
const auto* haystack_end = haystack + haystack_offsets[j] - prev_haystack_offset;
searcher.searchOneAll(haystack, haystack_end, &vec_res[from], res_callback);
prev_haystack_offset = haystack_offsets[j];
}
}
size_t accum = needles_size;
for (size_t i = 0; i < haystack_size; ++i) {
offsets_res[i] = accum;
accum += needles_size;
}
return Status::OK();
}
static Status vector_vector(const ColumnString::Chars& haystack_data,
const ColumnString::Offsets& haystack_offsets,
const IColumn& needles_data,
const ColumnArray::Offsets64& needles_offsets,
PaddedPODArray<Int32>& vec_res,
PaddedPODArray<UInt64>& offsets_res) {
size_t prev_haystack_offset = 0;
size_t prev_needles_offset = 0;
auto res_callback = [](const UInt8* start, const UInt8* end) -> Int32 {
return 1 + Impl::count_chars(reinterpret_cast<const char*>(start),
reinterpret_cast<const char*>(end));
};
offsets_res.reserve(haystack_offsets.size());
auto& nested_column =
vectorized::check_and_get_column<vectorized::ColumnNullable>(needles_data)
->get_nested_column();
const ColumnString* needles_data_string = check_and_get_column<ColumnString>(nested_column);
std::vector<StringRef> needles;
for (size_t i = 0; i < haystack_offsets.size(); ++i) {
needles.reserve(needles_offsets[i] - prev_needles_offset);
for (size_t j = prev_needles_offset; j < needles_offsets[i]; ++j) {
needles.emplace_back(needles_data_string->get_data_at(j));
}
const size_t needles_size = needles.size();
if (needles_size > std::numeric_limits<UInt8>::max())
return Status::InvalidArgument(
"number of arguments for function {} doesn't match: "
"passed {}, should be at most 255",
name, needles_size);
vec_res.resize(vec_res.size() + needles_size);
auto searcher = Impl::create_multi_searcher(needles);
std::fill(vec_res.begin() + vec_res.size() - needles_size, vec_res.end(), 0);
while (searcher.hasMoreToSearch()) {
const auto* haystack = &haystack_data[prev_haystack_offset];
const auto* haystack_end = haystack + haystack_offsets[i] - prev_haystack_offset;
searcher.searchOneAll(haystack, haystack_end,
&vec_res[vec_res.size() - needles_size], res_callback);
}
if (offsets_res.empty())
offsets_res.push_back(needles_size);
else
offsets_res.push_back(offsets_res.back() + needles_size);
prev_haystack_offset = haystack_offsets[i];
prev_needles_offset = needles_offsets[i];
needles.clear();
}
return Status::OK();
}
};
struct MultiSearcherImpl {
using MultiSearcher = MultiVolnitsky;
static MultiSearcher create_multi_searcher(const std::vector<StringRef>& needles) {
return MultiSearcher(needles);
}
static size_t count_chars(const char* begin, const char* end) { return end - begin; }
};
using FunctionMultiSearchAllPositions =
FunctionMultiStringPosition<FunctionMultiSearchAllPositionsImpl<MultiSearcherImpl>>;
void register_function_multi_string_position(SimpleFunctionFactory& factory) {
factory.register_function<FunctionMultiSearchAllPositions>();
}
} // namespace doris::vectorized