ClickHouse/dbms/include/DB/Functions/FunctionsStringArray.h
2017-03-11 04:15:25 +03:00

555 lines
17 KiB
C++
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#pragma once
#include <DB/DataTypes/DataTypeArray.h>
#include <DB/Columns/ColumnString.h>
#include <DB/Columns/ColumnFixedString.h>
#include <DB/Columns/ColumnConst.h>
#include <DB/Columns/ColumnArray.h>
#include <DB/Common/StringUtils.h>
#include <DB/Functions/IFunction.h>
#include <DB/Functions/Regexps.h>
#include <DB/DataTypes/DataTypeString.h>
#include <DB/IO/WriteHelpers.h>
namespace DB
{
/** Функции, разделяющие строки на массив строк или наоборот.
*
* splitByChar(sep, s)
* splitByString(sep, s)
* splitByRegexp(regexp, s)
*
* extractAll(s, regexp) - выделить из строки подпоследовательности, соответствующие регекспу.
* - первый subpattern, если в regexp-е есть subpattern;
* - нулевой subpattern (сматчившуюся часть, иначе);
* - инача, пустой массив
*
* arrayStringConcat(arr)
* arrayStringConcat(arr, delimiter)
* - склеить массив строк в одну строку через разделитель.
*
* alphaTokens(s) - выделить из строки подпоследовательности [a-zA-Z]+.
*
* Функции работы с URL расположены отдельно.
*/
using Pos = const char *;
/// Генераторы подстрок. Все они обладают общим интерфейсом.
class AlphaTokensImpl
{
private:
Pos pos;
Pos end;
public:
/// Получить имя фукнции.
static constexpr auto name = "alphaTokens";
static String getName() { return name; }
static size_t getNumberOfArguments() { return 1; }
/// Проверить типы агрументов функции.
static void checkArguments(const DataTypes & arguments)
{
if (!typeid_cast<const DataTypeString *>(&*arguments[0]))
throw Exception("Illegal type " + arguments[0]->getName() + " of first argument of function " + getName() + ". Must be String.",
ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
}
/// Инициализировать по аргументам функции.
void init(Block & block, const ColumnNumbers & arguments) {}
/// Вызывается для каждой следующей строки.
void set(Pos pos_, Pos end_)
{
pos = pos_;
end = end_;
}
/// Возвращает позицию аргумента, являющегося столбцом строк
size_t getStringsArgumentPosition()
{
return 0;
}
/// Получить следующий токен, если есть, или вернуть false.
bool get(Pos & token_begin, Pos & token_end)
{
/// Skip garbage
while (pos < end && !isAlphaASCII(*pos))
++pos;
if (pos == end)
return false;
token_begin = pos;
while (pos < end && isAlphaASCII(*pos))
++pos;
token_end = pos;
return true;
}
};
class SplitByCharImpl
{
private:
Pos pos;
Pos end;
char sep;
public:
static constexpr auto name = "splitByChar";
static String getName() { return name; }
static size_t getNumberOfArguments() { return 2; }
static void checkArguments(const DataTypes & arguments)
{
if (!typeid_cast<const DataTypeString *>(&*arguments[0]))
throw Exception("Illegal type " + arguments[0]->getName() + " of first argument of function " + getName() + ". Must be String.",
ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
if (!typeid_cast<const DataTypeString *>(&*arguments[1]))
throw Exception("Illegal type " + arguments[1]->getName() + " of second argument of function " + getName() + ". Must be String.",
ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
}
void init(Block & block, const ColumnNumbers & arguments)
{
const ColumnConstString * col = typeid_cast<const ColumnConstString *>(block.safeGetByPosition(arguments[0]).column.get());
if (!col)
throw Exception("Illegal column " + block.safeGetByPosition(arguments[0]).column->getName()
+ " of first argument of function " + getName() + ". Must be constant string.",
ErrorCodes::ILLEGAL_COLUMN);
const String & sep_str = col->getData();
if (sep_str.size() != 1)
throw Exception("Illegal separator for function " + getName() + ". Must be exactly one byte.");
sep = sep_str[0];
}
/// Возвращает позицию аргумента, являющегося столбцом строк
size_t getStringsArgumentPosition()
{
return 1;
}
void set(Pos pos_, Pos end_)
{
pos = pos_;
end = end_;
}
bool get(Pos & token_begin, Pos & token_end)
{
if (!pos)
return false;
token_begin = pos;
pos = reinterpret_cast<Pos>(memchr(pos, sep, end - pos));
if (pos)
{
token_end = pos;
++pos;
}
else
token_end = end;
return true;
}
};
class SplitByStringImpl
{
private:
Pos pos;
Pos end;
String sep;
public:
static constexpr auto name = "splitByString";
static String getName() { return name; }
static size_t getNumberOfArguments() { return 2; }
static void checkArguments(const DataTypes & arguments)
{
SplitByCharImpl::checkArguments(arguments);
}
void init(Block & block, const ColumnNumbers & arguments)
{
const ColumnConstString * col = typeid_cast<const ColumnConstString *>(block.safeGetByPosition(arguments[0]).column.get());
if (!col)
throw Exception("Illegal column " + block.safeGetByPosition(arguments[0]).column->getName()
+ " of first argument of function " + getName() + ". Must be constant string.",
ErrorCodes::ILLEGAL_COLUMN);
sep = col->getData();
}
/// Возвращает позицию аргумента, являющегося столбцом строк
size_t getStringsArgumentPosition()
{
return 1;
}
/// Вызывается для каждой следующей строки.
void set(Pos pos_, Pos end_)
{
pos = pos_;
end = end_;
}
/// Получить следующий токен, если есть, или вернуть false.
bool get(Pos & token_begin, Pos & token_end)
{
if (!pos)
return false;
token_begin = pos;
pos = reinterpret_cast<Pos>(memmem(pos, end - pos, sep.data(), sep.size()));
if (pos)
{
token_end = pos;
pos += sep.size();
}
else
token_end = end;
return true;
}
};
class ExtractAllImpl
{
private:
Regexps::Pool::Pointer re;
OptimizedRegularExpression::MatchVec matches;
size_t capture;
Pos pos;
Pos end;
public:
static constexpr auto name = "extractAll";
static String getName() { return name; }
static size_t getNumberOfArguments() { return 2; }
/// Проверить типы агрументов функции.
static void checkArguments( const DataTypes & arguments )
{
SplitByStringImpl::checkArguments(arguments);
}
/// Инициализировать по аргументам функции.
void init(Block & block, const ColumnNumbers & arguments)
{
const ColumnConstString * col = typeid_cast<const ColumnConstString *>(block.safeGetByPosition(arguments[1]).column.get());
if (!col)
throw Exception("Illegal column " + block.safeGetByPosition(arguments[1]).column->getName()
+ " of first argument of function " + getName() + ". Must be constant string.",
ErrorCodes::ILLEGAL_COLUMN);
re = Regexps::get<false, false>(col->getData());
capture = re->getNumberOfSubpatterns() > 0 ? 1 : 0;
matches.resize(capture + 1);
}
/// Возвращает позицию аргумента, являющегося столбцом строк
size_t getStringsArgumentPosition()
{
return 0;
}
/// Вызывается для каждой следующей строки.
void set(Pos pos_, Pos end_)
{
pos = pos_;
end = end_;
}
/// Получить следующий токен, если есть, или вернуть false.
bool get(Pos & token_begin, Pos & token_end)
{
if (!pos || pos > end)
return false;
if (!re->match(pos, end - pos, matches) || !matches[capture].length)
return false;
token_begin = pos + matches[capture].offset;
token_end = token_begin + matches[capture].length;
pos += matches[capture].offset + matches[capture].length;
return true;
}
};
/// Функция, принимающая строку, и возвращающая массив подстрок, создаваемый некоторым генератором.
template <typename Generator>
class FunctionTokens : public IFunction
{
public:
static constexpr auto name = Generator::name;
static FunctionPtr create(const Context & context) { return std::make_shared<FunctionTokens>(); }
String getName() const override
{
return name;
}
size_t getNumberOfArguments() const override { return Generator::getNumberOfArguments(); }
DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
{
Generator::checkArguments(arguments);
return std::make_shared<DataTypeArray>(std::make_shared<DataTypeString>());
}
void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result) override
{
Generator generator;
generator.init(block, arguments);
size_t arrayArgumentPosition = arguments[generator.getStringsArgumentPosition()];
const ColumnString * col_str = typeid_cast<const ColumnString *>(block.safeGetByPosition(arrayArgumentPosition).column.get());
const ColumnConstString * col_const_str =
typeid_cast<const ColumnConstString *>(block.safeGetByPosition(arrayArgumentPosition).column.get());
auto col_res = std::make_shared<ColumnArray>(std::make_shared<ColumnString>());
ColumnPtr col_res_holder = col_res;
ColumnString & res_strings = typeid_cast<ColumnString &>(col_res->getData());
ColumnArray::Offsets_t & res_offsets = col_res->getOffsets();
ColumnString::Chars_t & res_strings_chars = res_strings.getChars();
ColumnString::Offsets_t & res_strings_offsets = res_strings.getOffsets();
if (col_str)
{
const ColumnString::Chars_t & src_chars = col_str->getChars();
const ColumnString::Offsets_t & src_offsets = col_str->getOffsets();
res_offsets.reserve(src_offsets.size());
res_strings_offsets.reserve(src_offsets.size() * 5); /// Константа 5 - наугад.
res_strings_chars.reserve(src_chars.size());
Pos token_begin = nullptr;
Pos token_end = nullptr;
size_t size = src_offsets.size();
ColumnString::Offset_t current_src_offset = 0;
ColumnArray::Offset_t current_dst_offset = 0;
ColumnString::Offset_t current_dst_strings_offset = 0;
for (size_t i = 0; i < size; ++i)
{
Pos pos = reinterpret_cast<Pos>(&src_chars[current_src_offset]);
current_src_offset = src_offsets[i];
Pos end = reinterpret_cast<Pos>(&src_chars[current_src_offset]) - 1;
generator.set(pos, end);
size_t j = 0;
while (generator.get(token_begin, token_end))
{
size_t token_size = token_end - token_begin;
res_strings_chars.resize(res_strings_chars.size() + token_size + 1);
memcpySmallAllowReadWriteOverflow15(&res_strings_chars[current_dst_strings_offset], token_begin, token_size);
res_strings_chars[current_dst_strings_offset + token_size] = 0;
current_dst_strings_offset += token_size + 1;
res_strings_offsets.push_back(current_dst_strings_offset);
++j;
}
current_dst_offset += j;
res_offsets.push_back(current_dst_offset);
}
block.safeGetByPosition(result).column = col_res_holder;
}
else if (col_const_str)
{
String src = col_const_str->getData();
Array dst;
generator.set(src.data(), src.data() + src.size());
Pos token_begin = nullptr;
Pos token_end = nullptr;
while (generator.get(token_begin, token_end))
dst.push_back(String(token_begin, token_end - token_begin));
block.safeGetByPosition(result).column = std::make_shared<ColumnConstArray>(col_const_str->size(), dst, std::make_shared<DataTypeArray>(std::make_shared<DataTypeString>()));
}
else
throw Exception("Illegal columns " + block.safeGetByPosition(arrayArgumentPosition).column->getName()
+ ", " + block.safeGetByPosition(arrayArgumentPosition).column->getName()
+ " of arguments of function " + getName(),
ErrorCodes::ILLEGAL_COLUMN);
}
};
/// Склеивает массив строк в одну строку через разделитель.
class FunctionArrayStringConcat : public IFunction
{
private:
void executeInternal(
const ColumnString::Chars_t & src_chars,
const ColumnString::Offsets_t & src_string_offsets,
const ColumnArray::Offsets_t & src_array_offsets,
const char * delimiter, const size_t delimiter_size,
ColumnString::Chars_t & dst_chars,
ColumnString::Offsets_t & dst_string_offsets)
{
size_t size = src_array_offsets.size();
if (!size)
return;
/// С небольшим запасом - как будто разделитель идёт и после последней строки массива.
dst_chars.resize(
src_chars.size()
+ delimiter_size * src_string_offsets.size() /// Разделители после каждой строки...
+ src_array_offsets.size() /// Нулевой байт после каждой склеенной строки
- src_string_offsets.size()); /// Бывший нулевой байт после каждой строки массива
/// Будет столько строк, сколько было массивов.
dst_string_offsets.resize(src_array_offsets.size());
ColumnArray::Offset_t current_src_array_offset = 0;
ColumnString::Offset_t current_src_string_offset = 0;
ColumnString::Offset_t current_dst_string_offset = 0;
/// Цикл по массивам строк.
for (size_t i = 0; i < size; ++i)
{
/// Цикл по строкам внутри массива. /// NOTE Можно всё сделать за одно копирование, если разделитель имеет размер 1.
for (auto next_src_array_offset = src_array_offsets[i]; current_src_array_offset < next_src_array_offset; ++current_src_array_offset)
{
size_t bytes_to_copy = src_string_offsets[current_src_array_offset] - current_src_string_offset - 1;
memcpySmallAllowReadWriteOverflow15(
&dst_chars[current_dst_string_offset], &src_chars[current_src_string_offset], bytes_to_copy);
current_src_string_offset = src_string_offsets[current_src_array_offset];
current_dst_string_offset += bytes_to_copy;
if (current_src_array_offset + 1 != next_src_array_offset)
{
memcpy(&dst_chars[current_dst_string_offset], delimiter, delimiter_size);
current_dst_string_offset += delimiter_size;
}
}
dst_chars[current_dst_string_offset] = 0;
++current_dst_string_offset;
dst_string_offsets[i] = current_dst_string_offset;
}
dst_chars.resize(dst_string_offsets.back());
}
public:
static constexpr auto name = "arrayStringConcat";
static FunctionPtr create(const Context & context) { return std::make_shared<FunctionArrayStringConcat>(); }
String getName() const override
{
return name;
}
bool isVariadic() const override { return true; }
size_t getNumberOfArguments() const override { return 0; }
DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
{
if (arguments.size() != 1 && arguments.size() != 2)
throw Exception("Number of arguments for function " + getName() + " doesn't match: passed "
+ toString(arguments.size()) + ", should be 1 or 2.",
ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH);
const DataTypeArray * array_type = typeid_cast<const DataTypeArray *>(arguments[0].get());
if (!array_type || !typeid_cast<const DataTypeString *>(array_type->getNestedType().get()))
throw Exception("First argument for function " + getName() + " must be array of strings.", ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
if (arguments.size() == 2
&& !typeid_cast<const DataTypeString *>(arguments[1].get()))
throw Exception("Second argument for function " + getName() + " must be constant string.", ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
return std::make_shared<DataTypeString>();
}
void executeImpl(Block & block, const ColumnNumbers & arguments, size_t result) override
{
String delimiter;
if (arguments.size() == 2)
{
const ColumnConstString * col_delim = typeid_cast<const ColumnConstString *>(block.safeGetByPosition(arguments[1]).column.get());
if (!col_delim)
throw Exception("Second argument for function " + getName() + " must be constant string.", ErrorCodes::ILLEGAL_COLUMN);
delimiter = col_delim->getData();
}
if (const ColumnConstArray * col_const_arr = typeid_cast<const ColumnConstArray *>(block.safeGetByPosition(arguments[0]).column.get()))
{
auto col_res = std::make_shared<ColumnConstString>(col_const_arr->size(), "");
block.safeGetByPosition(result).column = col_res;
const Array & src_arr = col_const_arr->getData();
String & dst_str = col_res->getData();
for (size_t i = 0, size = src_arr.size(); i < size; ++i)
{
if (i != 0)
dst_str += delimiter;
dst_str += src_arr[i].get<const String &>();
}
}
else
{
const ColumnArray & col_arr = static_cast<const ColumnArray &>(*block.safeGetByPosition(arguments[0]).column);
const ColumnString & col_string = static_cast<const ColumnString &>(col_arr.getData());
std::shared_ptr<ColumnString> col_res = std::make_shared<ColumnString>();
block.safeGetByPosition(result).column = col_res;
executeInternal(
col_string.getChars(), col_string.getOffsets(), col_arr.getOffsets(),
delimiter.data(), delimiter.size(),
col_res->getChars(), col_res->getOffsets());
}
}
};
using FunctionAlphaTokens = FunctionTokens<AlphaTokensImpl> ;
using FunctionSplitByChar = FunctionTokens<SplitByCharImpl> ;
using FunctionSplitByString = FunctionTokens<SplitByStringImpl>;
using FunctionExtractAll = FunctionTokens<ExtractAllImpl> ;
}