ClickHouse/src/Functions/bitSlice.cpp
Alexander Tokmakov 70d1adfe4b
Better formatting for exception messages (#45449)
* save format string for NetException

* format exceptions

* format exceptions 2

* format exceptions 3

* format exceptions 4

* format exceptions 5

* format exceptions 6

* fix

* format exceptions 7

* format exceptions 8

* Update MergeTreeIndexGin.cpp

* Update AggregateFunctionMap.cpp

* Update AggregateFunctionMap.cpp

* fix
2023-01-24 00:13:58 +03:00

425 lines
17 KiB
C++

#include <Columns/ColumnConst.h>
#include <Columns/ColumnFixedString.h>
#include <Columns/ColumnString.h>
#include <DataTypes/DataTypeString.h>
#include <Functions/FunctionFactory.h>
#include <Functions/FunctionHelpers.h>
#include <Functions/GatherUtils/Algorithms.h>
#include <Functions/GatherUtils/Sinks.h>
#include <Functions/GatherUtils/Slices.h>
#include <Functions/GatherUtils/Sources.h>
#include <Functions/IFunction.h>
#include <IO/WriteHelpers.h>
namespace DB
{
using namespace GatherUtils;
namespace ErrorCodes
{
extern const int ILLEGAL_COLUMN;
extern const int ILLEGAL_TYPE_OF_ARGUMENT;
extern const int ZERO_ARRAY_OR_TUPLE_INDEX;
extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
}
class FunctionBitSlice : public IFunction
{
const UInt8 word_size = 8;
public:
static constexpr auto name = "bitSlice";
static FunctionPtr create(ContextPtr) { return std::make_shared<FunctionBitSlice>(); }
String getName() const override { return name; }
bool isVariadic() const override { return true; }
size_t getNumberOfArguments() const override { return 0; }
bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return true; }
bool useDefaultImplementationForConstants() const override { return true; }
DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
{
const size_t number_of_arguments = arguments.size();
if (number_of_arguments < 2 || number_of_arguments > 3)
throw Exception(ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH,
"Number of arguments for function {} doesn't match: passed {}, should be 2 or 3",
getName(), number_of_arguments);
if (!isString(arguments[0]) && !isStringOrFixedString(arguments[0]))
throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Illegal type {} of argument of function {}",
arguments[0]->getName(), getName());
if (arguments[0]->onlyNull())
return arguments[0];
if (!isNativeNumber(arguments[1]))
throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Illegal type {} of second argument of function {}",
arguments[1]->getName(), getName());
if (number_of_arguments == 3 && !isNativeNumber(arguments[2]))
throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Illegal type {} of second argument of function {}",
arguments[2]->getName(), getName());
return std::make_shared<DataTypeString>();
}
ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
{
size_t number_of_arguments = arguments.size();
ColumnPtr column_string = arguments[0].column;
ColumnPtr column_start = arguments[1].column;
ColumnPtr column_length;
std::optional<Int64> start_const;
std::optional<Int64> length_const;
if (const auto * column_start_const = checkAndGetColumn<ColumnConst>(column_start.get()))
{
start_const = column_start_const->getInt(0);
}
if (number_of_arguments == 3)
{
column_length = arguments[2].column;
if (const auto * column_length_const = checkAndGetColumn<ColumnConst>(column_length.get()))
length_const = column_length_const->getInt(0);
}
if (const ColumnString * col = checkAndGetColumn<ColumnString>(column_string.get()))
return executeForSource(column_start, column_length, start_const, length_const, StringSource(*col), input_rows_count);
else if (const ColumnFixedString * col_fixed = checkAndGetColumn<ColumnFixedString>(column_string.get()))
return executeForSource(
column_start, column_length, start_const, length_const, FixedStringSource(*col_fixed), input_rows_count);
else if (const ColumnConst * col_const = checkAndGetColumnConst<ColumnString>(column_string.get()))
return executeForSource(
column_start, column_length, start_const, length_const, ConstSource<StringSource>(*col_const), input_rows_count);
else if (const ColumnConst * col_const_fixed = checkAndGetColumnConst<ColumnFixedString>(column_string.get()))
return executeForSource(
column_start, column_length, start_const, length_const, ConstSource<FixedStringSource>(*col_const_fixed), input_rows_count);
else
throw Exception(ErrorCodes::ILLEGAL_COLUMN, "Illegal column {} of first argument of function {}",
arguments[0].column->getName(), getName());
}
template <class Source>
ColumnPtr executeForSource(
const ColumnPtr & column_start,
const ColumnPtr & column_length,
std::optional<Int64> start_const,
std::optional<Int64> length_const,
Source && source,
size_t input_rows_count) const
{
auto col_res = ColumnString::create();
if (!column_length)
{
if (start_const)
{
Int64 start_value = start_const.value();
if (start_value > 0)
bitSliceFromLeftConstantOffsetUnbounded(
source, StringSink(*col_res, input_rows_count), static_cast<size_t>(start_value - 1));
else if (start_value < 0)
bitSliceFromRightConstantOffsetUnbounded(
source, StringSink(*col_res, input_rows_count), -static_cast<size_t>(start_value));
else
throw Exception(ErrorCodes::ZERO_ARRAY_OR_TUPLE_INDEX, "Indices in strings are 1-based");
}
else
bitSliceDynamicOffsetUnbounded(source, StringSink(*col_res, input_rows_count), *column_start);
}
else
{
if (start_const && length_const)
{
Int64 start_value = start_const.value();
Int64 length_value = length_const.value();
if (start_value > 0)
bitSliceFromLeftConstantOffsetBounded(
source, StringSink(*col_res, input_rows_count), static_cast<size_t>(start_value - 1), length_value);
else if (start_value < 0)
bitSliceFromRightConstantOffsetBounded(
source, StringSink(*col_res, input_rows_count), -static_cast<size_t>(start_value), length_value);
else
throw Exception(ErrorCodes::ZERO_ARRAY_OR_TUPLE_INDEX, "Indices in strings are 1-based");
}
else
bitSliceDynamicOffsetBounded(source, StringSink(*col_res, input_rows_count), *column_start, *column_length);
}
return col_res;
}
void writeSliceWithLeftShift(const StringSource::Slice & slice, StringSink & sink, size_t shift_bit, size_t abandon_last_bit = 0) const
{
if (!shift_bit && !abandon_last_bit)
{
writeSlice(slice, sink);
return;
}
size_t size = slice.size;
if (!size)
return;
bool abandon_last_byte = abandon_last_bit + shift_bit >= word_size;
if (abandon_last_byte) // shift may eliminate last byte
size--;
sink.elements.resize(sink.current_offset + size);
UInt8 * out = &sink.elements[sink.current_offset];
const UInt8 * input = slice.data;
for (size_t i = 0; i < size - 1; i++)
{
out[i] = (input[i] << shift_bit) | (input[i + 1] >> (word_size - shift_bit));
}
if (abandon_last_byte)
{
out[size - 1] = (input[size - 1] << shift_bit) | (input[size] >> (word_size - shift_bit));
out[size - 1] = out[size - 1] & (0xFF << (abandon_last_bit + shift_bit - word_size));
}
else
{
out[size - 1] = (input[size - 1] << shift_bit) & (0xFF << (abandon_last_bit + shift_bit));
}
sink.current_offset += size;
}
template <class Source>
void bitSliceFromLeftConstantOffsetUnbounded(Source && src, StringSink && sink, size_t offset) const
{
size_t offset_byte = offset / word_size;
size_t offset_bit = offset % word_size;
while (!src.isEnd())
{
auto sl = src.getSliceFromLeft(offset_byte);
if (sl.size)
writeSliceWithLeftShift(sl, sink, offset_bit);
sink.next();
src.next();
}
}
template <class Source>
void bitSliceFromRightConstantOffsetUnbounded(Source && src, StringSink && sink, size_t offset) const
{
size_t offset_byte = offset / word_size;
size_t offset_bit = (word_size - (offset % word_size)) % word_size; // offset_bit always represent left offset bit
if (offset_bit)
offset_byte++;
while (!src.isEnd())
{
auto slice = src.getSliceFromRight(offset_byte);
size_t size = src.getElementSize();
bool left_truncate = offset_byte > size;
size_t shift_bit = left_truncate ? 0 : offset_bit;
if (slice.size)
writeSliceWithLeftShift(slice, sink, shift_bit);
sink.next();
src.next();
}
}
template <class Source>
void bitSliceDynamicOffsetUnbounded(Source && src, StringSink && sink, const IColumn & offset_column) const
{
while (!src.isEnd())
{
auto row_num = src.rowNum();
Int64 start = offset_column.getInt(row_num);
if (start != 0)
{
typename std::decay_t<Source>::Slice slice;
size_t shift_bit;
if (start > 0)
{
UInt64 offset = start - 1;
size_t offset_byte = offset / word_size;
size_t offset_bit = offset % word_size;
shift_bit = offset_bit;
slice = src.getSliceFromLeft(offset_byte);
}
else
{
UInt64 offset = -static_cast<UInt64>(start);
size_t offset_byte = offset / word_size;
size_t offset_bit = (word_size - (offset % word_size)) % word_size; // offset_bit always represent left offset bit
if (offset_bit)
offset_byte++;
size_t size = src.getElementSize();
bool left_truncate = offset_byte > size;
shift_bit = left_truncate ? 0 : offset_bit;
slice = src.getSliceFromRight(offset_byte);
}
if (slice.size)
writeSliceWithLeftShift(slice, sink, shift_bit);
}
sink.next();
src.next();
}
}
template <class Source>
void bitSliceFromLeftConstantOffsetBounded(Source && src, StringSink && sink, size_t offset, ssize_t length) const
{
size_t offset_byte = offset / word_size;
size_t offset_bit = offset % word_size;
size_t shift_bit = offset_bit;
size_t length_byte = 0;
size_t over_bit = 0;
if (length > 0)
{
length_byte = (length + offset_bit) / word_size;
over_bit = (length + offset_bit) % word_size;
if (over_bit && (length_byte || over_bit > offset_bit)) // begin and end are not in same byte OR there are gaps
length_byte++;
}
while (!src.isEnd())
{
ssize_t remain_byte = src.getElementSize() - offset_byte;
if (length < 0)
{
length_byte = std::max(remain_byte + (length / word_size), static_cast<ssize_t>(0));
over_bit = word_size + (length % word_size);
if (length_byte == 1 && over_bit <= offset_bit) // begin and end are in same byte AND there are no gaps
length_byte = 0;
}
bool right_truncate = static_cast<ssize_t>(length_byte) > remain_byte;
size_t abandon_last_bit = (over_bit && !right_truncate) ? word_size - over_bit : 0;
auto slice = src.getSliceFromLeft(offset_byte, length_byte);
if (slice.size)
writeSliceWithLeftShift(slice, sink, shift_bit, abandon_last_bit);
sink.next();
src.next();
}
}
template <class Source>
void bitSliceFromRightConstantOffsetBounded(Source && src, StringSink && sink, size_t offset, ssize_t length) const
{
size_t offset_byte = offset / word_size;
size_t offset_bit = (word_size - (offset % word_size)) % word_size; // offset_bit always represent left offset bit
if (offset_bit)
offset_byte++;
size_t length_byte = 0;
size_t over_bit = 0;
if (length > 0)
{
length_byte = (length + offset_bit) / word_size;
over_bit = (length + offset_bit) % word_size;
if (over_bit && (length_byte || over_bit > offset_bit)) // begin and end are not in same byte OR there are gaps
length_byte++;
}
while (!src.isEnd())
{
size_t size = src.getElementSize();
if (length < 0)
{
length_byte = std::max(static_cast<ssize_t>(offset_byte) + (length / word_size), static_cast<ssize_t>(0));
over_bit = word_size + (length % word_size);
if (length_byte == 1 && over_bit <= offset_bit) // begin and end are in same byte AND there are no gaps
length_byte = 0;
}
bool left_truncate = offset_byte > size;
bool right_truncate = length_byte > offset_byte;
size_t shift_bit = left_truncate ? 0 : offset_bit;
size_t abandon_last_bit = (over_bit && !right_truncate) ? word_size - over_bit : 0;
auto slice = src.getSliceFromRight(offset_byte, length_byte);
if (slice.size)
writeSliceWithLeftShift(slice, sink, shift_bit, abandon_last_bit);
sink.next();
src.next();
}
}
template <class Source>
void bitSliceDynamicOffsetBounded(Source && src, StringSink && sink, const IColumn & offset_column, const IColumn & length_column) const
{
while (!src.isEnd())
{
size_t row_num = src.rowNum();
Int64 start = offset_column.getInt(row_num);
Int64 length = length_column.getInt(row_num);
if (start && length)
{
bool left_offset = start > 0;
size_t offset = left_offset ? static_cast<size_t>(start - 1) : -static_cast<size_t>(start);
size_t size = src.getElementSize();
size_t offset_byte;
size_t offset_bit;
size_t shift_bit;
if (left_offset)
{
offset_byte = offset / word_size;
offset_bit = offset % word_size;
shift_bit = offset_bit;
}
else
{
offset_byte = offset / word_size;
offset_bit = (word_size - (offset % word_size)) % word_size; // offset_bit always represent left offset bit
if (offset_bit)
offset_byte++;
bool left_truncate = offset_byte > size;
shift_bit = left_truncate ? 0 : offset_bit;
}
ssize_t remain_byte = left_offset ? size - offset_byte : offset_byte;
size_t length_byte;
size_t over_bit;
if (length > 0)
{
length_byte = (length + offset_bit) / word_size;
over_bit = (length + offset_bit) % word_size;
if (over_bit && (length_byte || (over_bit > offset_bit))) // begin and end are not in same byte OR there are gaps
length_byte++;
}
else
{
length_byte = std::max(remain_byte + (static_cast<ssize_t>(length) / word_size), static_cast<ssize_t>(0));
over_bit = word_size + (length % word_size);
if (length_byte == 1 && over_bit <= offset_bit) // begin and end are in same byte AND there are no gaps
length_byte = 0;
}
bool right_truncate = static_cast<ssize_t>(length_byte) > remain_byte;
size_t abandon_last_bit = (over_bit && !right_truncate) ? word_size - over_bit : 0;
auto slice = left_offset ? src.getSliceFromLeft(offset_byte, length_byte) : src.getSliceFromRight(offset_byte, length_byte);
if (slice.size)
writeSliceWithLeftShift(slice, sink, shift_bit, abandon_last_bit);
}
sink.next();
src.next();
}
}
};
REGISTER_FUNCTION(BitSlice)
{
factory.registerFunction<FunctionBitSlice>();
}
}