ClickHouse

1
#include <memory>
2
#include <Columns/ColumnString.h>
3
#include <DataTypes/DataTypeString.h>
4
#include <Formats/FormatFactory.h>
5
#include <Functions/FunctionFactory.h>
6
#include <Functions/FunctionHelpers.h>
7
#include <Functions/IFunction.h>
8
#include <IO/WriteBufferFromVector.h>
9
#include <IO/WriteHelpers.h>
10
#include <Processors/Formats/IOutputFormat.h>
11
#include <Processors/Formats/IRowOutputFormat.h>
12
#include <base/map.h>
13

14

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

24
namespace
25
{
26

27
/** formatRow(<format>, x, y, ...) is a function that allows you to use RowOutputFormat over
28
  * several columns to generate a string per row, such as CSV, TSV, JSONEachRow, etc.
29
  * formatRowNoNewline(...) trims the newline character of each row.
30
  */
31
template <bool no_newline>
32
class FunctionFormatRow : public IFunction
33
{
34
public:
35
    static constexpr auto name = no_newline ? "formatRowNoNewline" : "formatRow";
36

37
    FunctionFormatRow(String format_name_, Names arguments_column_names_, ContextPtr context_)
38
        : format_name(std::move(format_name_))
39
        , arguments_column_names(std::move(arguments_column_names_))
40
        , context(std::move(context_))
41
    {
42
        FormatFactory::instance().checkFormatName(format_name);
43
    }
44

45
    String getName() const override { return name; }
46
    size_t getNumberOfArguments() const override { return 0; }
47
    bool useDefaultImplementationForNulls() const override { return false; }
48
    bool useDefaultImplementationForConstants() const override { return true; }
49
    bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return true; }
50
    ColumnNumbers getArgumentsThatAreAlwaysConstant() const override { return {0}; }
51

52
    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
53
    {
54
        auto col_str = ColumnString::create();
55
        ColumnString::Chars & vec = col_str->getChars();
56
        WriteBufferFromVector buffer(vec);
57
        ColumnString::Offsets & offsets = col_str->getOffsets();
58
        offsets.resize(input_rows_count);
59

60
        Block arg_columns;
61

62
        size_t arguments_size = arguments.size();
63
        for (size_t i = 1; i < arguments_size; ++i)
64
        {
65
            auto argument_column = arguments[i];
66
            argument_column.name = arguments_column_names[i];
67
            arg_columns.insert(std::move(argument_column));
68
        }
69

70
        materializeBlockInplace(arg_columns);
71
        auto format_settings = getFormatSettings(context);
72
        auto out = FormatFactory::instance().getOutputFormat(format_name, buffer, arg_columns, context, format_settings);
73

74
        /// This function make sense only for row output formats.
75
        auto * row_output_format = dynamic_cast<IRowOutputFormat *>(out.get());
76
        if (!row_output_format)
77
            throw Exception(ErrorCodes::BAD_ARGUMENTS,
78
                            "Cannot turn rows into a {} format strings. {} function supports only row output formats",
79
                            format_name, getName());
80

81
        auto columns = arg_columns.getColumns();
82
        for (size_t i = 0; i != input_rows_count; ++i)
83
        {
84
            row_output_format->writePrefixIfNeeded();
85
            row_output_format->writeRow(columns, i);
86
            row_output_format->finalize();
87
            if constexpr (no_newline)
88
            {
89
                // replace '\n' with '\0'
90
                if (buffer.position() != buffer.buffer().begin() && buffer.position()[-1] == '\n')
91
                    buffer.position()[-1] = '\0';
92
            }
93
            else
94
                writeChar('\0', buffer);
95

96
            offsets[i] = buffer.count();
97
            row_output_format->resetFormatter();
98
        }
99

100
        return col_str;
101
    }
102

103
private:
104
    String format_name;
105
    Names arguments_column_names;
106
    ContextPtr context;
107
};
108

109
template <bool no_newline>
110
class FormatRowOverloadResolver : public IFunctionOverloadResolver
111
{
112
public:
113
    static constexpr auto name = no_newline ? "formatRowNoNewline" : "formatRow";
114
    static FunctionOverloadResolverPtr create(ContextPtr context) { return std::make_unique<FormatRowOverloadResolver>(context); }
115
    explicit FormatRowOverloadResolver(ContextPtr context_) : context(context_) { }
116
    String getName() const override { return name; }
117
    bool isVariadic() const override { return true; }
118
    size_t getNumberOfArguments() const override { return 0; }
119
    ColumnNumbers getArgumentsThatAreAlwaysConstant() const override { return {0}; }
120
    bool useDefaultImplementationForNulls() const override { return false; }
121

122
    FunctionBasePtr buildImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr & return_type) const override
123
    {
124
        if (arguments.size() < 2)
125
            throw Exception(ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH,
126
                "Function {} requires at least two arguments: the format name and its output expression(s)", getName());
127

128
        Names arguments_column_names;
129
        arguments_column_names.reserve(arguments.size());
130
        for (const auto & argument : arguments)
131
            arguments_column_names.push_back(argument.name);
132

133
        if (const auto * name_col = checkAndGetColumnConst<ColumnString>(arguments.at(0).column.get()))
134
            return std::make_unique<FunctionToFunctionBaseAdaptor>(
135
                std::make_shared<FunctionFormatRow<no_newline>>(name_col->getValue<String>(), std::move(arguments_column_names), context),
136
                collections::map<DataTypes>(arguments, [](const auto & elem) { return elem.type; }),
137
                return_type);
138
        else
139
            throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "First argument to {} must be a format name", getName());
140
    }
141

142
    DataTypePtr getReturnTypeImpl(const DataTypes &) const override { return std::make_shared<DataTypeString>(); }
143

144
private:
145
    ContextPtr context;
146
};
147

148
}
149

150
REGISTER_FUNCTION(FormatRow)
151
{
152
    factory.registerFunction<FormatRowOverloadResolver<true>>();
153
    factory.registerFunction<FormatRowOverloadResolver<false>>();
154
}
155

156
}
157