ClickHouse/src/Functions/tupleElement.cpp
2021-08-09 18:01:08 +03:00

153 lines
5.0 KiB
C++

#include <Functions/IFunction.h>
#include <Functions/FunctionFactory.h>
#include <Functions/FunctionHelpers.h>
#include <DataTypes/IDataType.h>
#include <DataTypes/DataTypeTuple.h>
#include <DataTypes/DataTypeArray.h>
#include <Columns/ColumnTuple.h>
#include <Columns/ColumnArray.h>
#include <Columns/ColumnString.h>
#include <Columns/ColumnsNumber.h>
#include <Common/assert_cast.h>
#include <memory>
namespace DB
{
namespace ErrorCodes
{
extern const int ILLEGAL_TYPE_OF_ARGUMENT;
extern const int ILLEGAL_INDEX;
}
namespace
{
/** Extract element of tuple by constant index or name. The operation is essentially free.
* Also the function looks through Arrays: you can get Array of tuple elements from Array of Tuples.
*/
class FunctionTupleElement : public IFunction
{
public:
static constexpr auto name = "tupleElement";
static FunctionPtr create(ContextPtr)
{
return std::make_shared<FunctionTupleElement>();
}
String getName() const override
{
return name;
}
size_t getNumberOfArguments() const override
{
return 2;
}
bool useDefaultImplementationForConstants() const override
{
return true;
}
ColumnNumbers getArgumentsThatAreAlwaysConstant() const override
{
return {1};
}
bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return false; }
DataTypePtr getReturnTypeImpl(const ColumnsWithTypeAndName & arguments) const override
{
size_t count_arrays = 0;
const IDataType * tuple_col = arguments[0].type.get();
while (const DataTypeArray * array = checkAndGetDataType<DataTypeArray>(tuple_col))
{
tuple_col = array->getNestedType().get();
++count_arrays;
}
const DataTypeTuple * tuple = checkAndGetDataType<DataTypeTuple>(tuple_col);
if (!tuple)
throw Exception("First argument for function " + getName() + " must be tuple or array of tuple.", ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
size_t index = getElementNum(arguments[1].column, *tuple);
DataTypePtr out_return_type = tuple->getElements()[index];
for (; count_arrays; --count_arrays)
out_return_type = std::make_shared<DataTypeArray>(out_return_type);
return out_return_type;
}
ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t /*input_rows_count*/) const override
{
Columns array_offsets;
const auto & first_arg = arguments[0];
const IDataType * tuple_type = first_arg.type.get();
const IColumn * tuple_col = first_arg.column.get();
while (const DataTypeArray * array_type = checkAndGetDataType<DataTypeArray>(tuple_type))
{
const ColumnArray * array_col = assert_cast<const ColumnArray *>(tuple_col);
tuple_type = array_type->getNestedType().get();
tuple_col = &array_col->getData();
array_offsets.push_back(array_col->getOffsetsPtr());
}
const DataTypeTuple * tuple_type_concrete = checkAndGetDataType<DataTypeTuple>(tuple_type);
const ColumnTuple * tuple_col_concrete = checkAndGetColumn<ColumnTuple>(tuple_col);
if (!tuple_type_concrete || !tuple_col_concrete)
throw Exception("First argument for function " + getName() + " must be tuple or array of tuple.", ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
size_t index = getElementNum(arguments[1].column, *tuple_type_concrete);
ColumnPtr res = tuple_col_concrete->getColumns()[index];
/// Wrap into Arrays
for (auto it = array_offsets.rbegin(); it != array_offsets.rend(); ++it)
res = ColumnArray::create(res, *it);
return res;
}
private:
size_t getElementNum(const ColumnPtr & index_column, const DataTypeTuple & tuple) const
{
if (
checkAndGetColumnConst<ColumnUInt8>(index_column.get())
|| checkAndGetColumnConst<ColumnUInt16>(index_column.get())
|| checkAndGetColumnConst<ColumnUInt32>(index_column.get())
|| checkAndGetColumnConst<ColumnUInt64>(index_column.get())
)
{
size_t index = index_column->getUInt(0);
if (index == 0)
throw Exception("Indices in tuples are 1-based.", ErrorCodes::ILLEGAL_INDEX);
if (index > tuple.getElements().size())
throw Exception("Index for tuple element is out of range.", ErrorCodes::ILLEGAL_INDEX);
return index - 1;
}
else if (const auto * name_col = checkAndGetColumnConst<ColumnString>(index_column.get()))
{
return tuple.getPositionByName(name_col->getValue<String>());
}
else
throw Exception("Second argument to " + getName() + " must be a constant UInt or String", ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT);
}
};
}
void registerFunctionTupleElement(FunctionFactory & factory)
{
factory.registerFunction<FunctionTupleElement>();
}
}