ClickHouse

1
#include <Columns/ColumnFixedString.h>
2
#include <Columns/ColumnString.h>
3
#include <Columns/ColumnStringHelpers.h>
4
#include <DataTypes/DataTypeString.h>
5
#include <Functions/FunctionFactory.h>
6
#include <Functions/FunctionHelpers.h>
7
#include <Functions/IFunction.h>
8
#include <Functions/formatString.h>
9
#include <IO/WriteHelpers.h>
10

11
#include <memory>
12
#include <string>
13
#include <vector>
14

15

16
namespace DB
17
{
18
namespace ErrorCodes
19
{
20
    extern const int ILLEGAL_COLUMN;
21
    extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH;
22
}
23

24
namespace
25
{
26

27
class FormatFunction : public IFunction
28
{
29
public:
30
    static constexpr auto name = "format";
31

32
    static FunctionPtr create(ContextPtr) { return std::make_shared<FormatFunction>(); }
33

34
    String getName() const override { return name; }
35

36
    bool isVariadic() const override { return true; }
37

38
    bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return true; }
39

40
    size_t getNumberOfArguments() const override { return 0; }
41

42
    bool useDefaultImplementationForConstants() const override { return true; }
43
    ColumnNumbers getArgumentsThatAreAlwaysConstant() const override { return {0}; }
44

45
    DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
46
    {
47
        if (arguments.size() < 2)
48
            throw Exception(
49
                ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH,
50
                "Number of arguments for function {} doesn't match: passed {}, should be at least 2",
51
                getName(),
52
                arguments.size());
53

54
        return std::make_shared<DataTypeString>();
55
    }
56

57
    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
58
    {
59
        const ColumnPtr & c0 = arguments[0].column;
60
        const ColumnConst * c0_const_string = typeid_cast<const ColumnConst *>(&*c0);
61

62
        if (!c0_const_string)
63
            throw Exception(ErrorCodes::ILLEGAL_COLUMN, "First argument of function {} must be constant string", getName());
64

65
        String pattern = c0_const_string->getValue<String>();
66

67
        auto col_res = ColumnString::create();
68

69
        std::vector<const ColumnString::Chars *> data(arguments.size() - 1);
70
        std::vector<const ColumnString::Offsets *> offsets(arguments.size() - 1);
71
        std::vector<size_t> fixed_string_sizes(arguments.size() - 1);
72
        std::vector<std::optional<String>> constant_strings(arguments.size() - 1);
73
        std::vector<ColumnString::MutablePtr> converted_col_ptrs(arguments.size() - 1);
74

75
        bool has_column_string = false;
76
        bool has_column_fixed_string = false;
77
        for (size_t i = 1; i < arguments.size(); ++i)
78
        {
79
            const ColumnPtr & column = arguments[i].column;
80
            if (const ColumnString * col = checkAndGetColumn<ColumnString>(column.get()))
81
            {
82
                has_column_string = true;
83
                data[i - 1] = &col->getChars();
84
                offsets[i - 1] = &col->getOffsets();
85
            }
86
            else if (const ColumnFixedString * fixed_col = checkAndGetColumn<ColumnFixedString>(column.get()))
87
            {
88
                has_column_fixed_string = true;
89
                data[i - 1] = &fixed_col->getChars();
90
                fixed_string_sizes[i - 1] = fixed_col->getN();
91
            }
92
            else if (const ColumnConst * const_col = checkAndGetColumnConstStringOrFixedString(column.get()))
93
            {
94
                constant_strings[i - 1] = const_col->getValue<String>();
95
            }
96
            else
97
            {
98
                /// A non-String/non-FixedString-type argument: use the default serialization to convert it to String
99
                auto full_column = column->convertToFullIfNeeded();
100
                auto serialization = arguments[i].type->getDefaultSerialization();
101
                auto converted_col_str = ColumnString::create();
102
                ColumnStringHelpers::WriteHelper write_helper(*converted_col_str, column->size());
103
                auto & write_buffer = write_helper.getWriteBuffer();
104
                FormatSettings format_settings;
105
                for (size_t row = 0; row < column->size(); ++row)
106
                {
107
                    serialization->serializeText(*full_column, row, write_buffer, format_settings);
108
                    write_helper.rowWritten();
109
                }
110
                write_helper.finalize();
111

112
                /// Keep the pointer alive
113
                converted_col_ptrs[i - 1] = std::move(converted_col_str);
114

115
                /// Same as the normal `ColumnString` branch
116
                has_column_string = true;
117
                data[i - 1] = &converted_col_ptrs[i - 1]->getChars();
118
                offsets[i - 1] = &converted_col_ptrs[i - 1]->getOffsets();
119
            }
120
        }
121

122
        FormatStringImpl::formatExecute(
123
            has_column_string,
124
            has_column_fixed_string,
125
            std::move(pattern),
126
            data,
127
            offsets,
128
            fixed_string_sizes,
129
            constant_strings,
130
            col_res->getChars(),
131
            col_res->getOffsets(),
132
            input_rows_count);
133

134
        return col_res;
135
    }
136
};
137

138

139
using FunctionFormat = FormatFunction;
140

141
}
142

143
REGISTER_FUNCTION(Format)
144
{
145
    factory.registerFunction<FunctionFormat>();
146
}
147

148
}
149