#pragma once
#include
#include
#include
#include
#include
#include
#include
#include
#include
#include
#include
#include
namespace DB
{
using namespace GatherUtils;
namespace ErrorCodes
{
extern const int ILLEGAL_COLUMN;
extern const int LOGICAL_ERROR;
extern const int ILLEGAL_TYPE_OF_ARGUMENT;
}
struct NameStartsWith
{
static constexpr auto name = "startsWith";
};
struct NameEndsWith
{
static constexpr auto name = "endsWith";
};
DECLARE_MULTITARGET_CODE(
template
class FunctionStartsEndsWith : public IFunction
{
public:
static constexpr auto name = Name::name;
String getName() const override
{
return name;
}
bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override
{
return true;
}
size_t getNumberOfArguments() const override
{
return 2;
}
bool useDefaultImplementationForConstants() const override
{
return true;
}
DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
{
if (isStringOrFixedString(arguments[0]) && isStringOrFixedString(arguments[1]))
return std::make_shared();
if (isArray(arguments[0]) && isArray(arguments[1]))
return std::make_shared();
throw Exception("Illegal type " + arguments[0]->getName() + " of argument of function " + getName(), ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
}
ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
{
auto data_type = arguments[0].type;
if (isStringOrFixedString(*data_type))
return executeImplString(arguments, {}, input_rows_count);
if (isArray(data_type))
return executeImplArray(arguments, {}, input_rows_count);
return {};
}
private:
ColumnPtr executeImplArray(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const
{
DataTypePtr common_type = getLeastSupertype(collections::map(arguments, [](auto & arg) { return arg.type; }));
Columns preprocessed_columns(2);
for (size_t i = 0; i < 2; ++i)
preprocessed_columns[i] = castColumn(arguments[i], common_type);
std::vector> sources;
for (auto & argument_column : preprocessed_columns)
{
bool is_const = false;
if (const auto * argument_column_const = typeid_cast(argument_column.get()))
{
is_const = true;
argument_column = argument_column_const->getDataColumnPtr();
}
if (const auto * argument_column_array = typeid_cast(argument_column.get()))
sources.emplace_back(GatherUtils::createArraySource(*argument_column_array, is_const, input_rows_count));
else
throw Exception{"Arguments for function " + getName() + " must be arrays.", ErrorCodes::LOGICAL_ERROR};
}
auto result_column = ColumnUInt8::create(input_rows_count);
auto * result_column_ptr = typeid_cast(result_column.get());
if constexpr (std::is_same_v)
GatherUtils::sliceHas(*sources[0], *sources[1], GatherUtils::ArraySearchType::StartsWith, *result_column_ptr);
else
GatherUtils::sliceHas(*sources[0], *sources[1], GatherUtils::ArraySearchType::EndsWith, *result_column_ptr);
return result_column;
}
ColumnPtr executeImplString(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const
{
const IColumn * haystack_column = arguments[0].column.get();
const IColumn * needle_column = arguments[1].column.get();
auto col_res = ColumnVector::create();
typename ColumnVector::Container & vec_res = col_res->getData();
vec_res.resize(input_rows_count);
if (const ColumnString * haystack = checkAndGetColumn(haystack_column))
dispatch(StringSource(*haystack), needle_column, vec_res);
else if (const ColumnFixedString * haystack_fixed = checkAndGetColumn(haystack_column))
dispatch(FixedStringSource(*haystack_fixed), needle_column, vec_res);
else if (const ColumnConst * haystack_const = checkAndGetColumnConst(haystack_column))
dispatch>(ConstSource(*haystack_const), needle_column, vec_res);
else if (const ColumnConst * haystack_const_fixed = checkAndGetColumnConst(haystack_column))
dispatch>(ConstSource(*haystack_const_fixed), needle_column, vec_res);
else
throw Exception("Illegal combination of columns as arguments of function " + getName(), ErrorCodes::ILLEGAL_COLUMN);
return col_res;
}
template
void dispatch(HaystackSource haystack_source, const IColumn * needle_column, PaddedPODArray & res_data) const
{
if (const ColumnString * needle = checkAndGetColumn(needle_column))
execute(haystack_source, StringSource(*needle), res_data);
else if (const ColumnFixedString * needle_fixed = checkAndGetColumn(needle_column))
execute(haystack_source, FixedStringSource(*needle_fixed), res_data);
else if (const ColumnConst * needle_const = checkAndGetColumnConst(needle_column))
execute>(haystack_source, ConstSource(*needle_const), res_data);
else if (const ColumnConst * needle_const_fixed = checkAndGetColumnConst(needle_column))
execute>(haystack_source, ConstSource(*needle_const_fixed), res_data);
else
throw Exception("Illegal combination of columns as arguments of function " + getName(), ErrorCodes::ILLEGAL_COLUMN);
}
template
static void execute(HaystackSource haystack_source, NeedleSource needle_source, PaddedPODArray & res_data)
{
size_t row_num = 0;
while (!haystack_source.isEnd())
{
auto haystack = haystack_source.getWhole();
auto needle = needle_source.getWhole();
if (needle.size > haystack.size)
{
res_data[row_num] = false;
}
else
{
if constexpr (std::is_same_v)
{
res_data[row_num] = StringRef(haystack.data, needle.size) == StringRef(needle.data, needle.size);
}
else /// endsWith
{
res_data[row_num] = StringRef(haystack.data + haystack.size - needle.size, needle.size) == StringRef(needle.data, needle.size);
}
}
haystack_source.next();
needle_source.next();
++row_num;
}
}
};
) // DECLARE_MULTITARGET_CODE
template
class FunctionStartsEndsWith : public TargetSpecific::Default::FunctionStartsEndsWith
{
public:
explicit FunctionStartsEndsWith(ContextPtr context) : selector(context)
{
selector.registerImplementation>();
#if USE_MULTITARGET_CODE
selector.registerImplementation>();
selector.registerImplementation>();
selector.registerImplementation>();
selector.registerImplementation>();
#endif
}
ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr & result_type, size_t input_rows_count) const override
{
return selector.selectAndExecute(arguments, result_type, input_rows_count);
}
static FunctionPtr create(ContextPtr context)
{
return std::make_shared>(context);
}
private:
ImplementationSelector selector;
};
}