1 | #include "duckdb/function/table/range.hpp" |
2 | #include "duckdb/common/algorithm.hpp" |
3 | |
4 | namespace duckdb { |
5 | |
6 | struct RepeatFunctionData : public TableFunctionData { |
7 | RepeatFunctionData(Value value, idx_t target_count) : value(std::move(value)), target_count(target_count) { |
8 | } |
9 | |
10 | Value value; |
11 | idx_t target_count; |
12 | }; |
13 | |
14 | struct RepeatOperatorData : public GlobalTableFunctionState { |
15 | RepeatOperatorData() : current_count(0) { |
16 | } |
17 | idx_t current_count; |
18 | }; |
19 | |
20 | static unique_ptr<FunctionData> RepeatBind(ClientContext &context, TableFunctionBindInput &input, |
21 | vector<LogicalType> &return_types, vector<string> &names) { |
22 | // the repeat function returns the type of the first argument |
23 | auto &inputs = input.inputs; |
24 | return_types.push_back(x: inputs[0].type()); |
25 | names.push_back(x: inputs[0].ToString()); |
26 | if (inputs[1].IsNull()) { |
27 | throw BinderException("Repeat second parameter cannot be NULL" ); |
28 | } |
29 | return make_uniq<RepeatFunctionData>(args&: inputs[0], args: inputs[1].GetValue<int64_t>()); |
30 | } |
31 | |
32 | static unique_ptr<GlobalTableFunctionState> RepeatInit(ClientContext &context, TableFunctionInitInput &input) { |
33 | return make_uniq<RepeatOperatorData>(); |
34 | } |
35 | |
36 | static void RepeatFunction(ClientContext &context, TableFunctionInput &data_p, DataChunk &output) { |
37 | auto &bind_data = data_p.bind_data->Cast<RepeatFunctionData>(); |
38 | auto &state = data_p.global_state->Cast<RepeatOperatorData>(); |
39 | |
40 | idx_t remaining = MinValue<idx_t>(a: bind_data.target_count - state.current_count, STANDARD_VECTOR_SIZE); |
41 | output.data[0].Reference(value: bind_data.value); |
42 | output.SetCardinality(remaining); |
43 | state.current_count += remaining; |
44 | } |
45 | |
46 | static unique_ptr<NodeStatistics> RepeatCardinality(ClientContext &context, const FunctionData *bind_data_p) { |
47 | auto &bind_data = bind_data_p->Cast<RepeatFunctionData>(); |
48 | return make_uniq<NodeStatistics>(args: bind_data.target_count, args: bind_data.target_count); |
49 | } |
50 | |
51 | void RepeatTableFunction::RegisterFunction(BuiltinFunctions &set) { |
52 | TableFunction repeat("repeat" , {LogicalType::ANY, LogicalType::BIGINT}, RepeatFunction, RepeatBind, RepeatInit); |
53 | repeat.cardinality = RepeatCardinality; |
54 | set.AddFunction(function: repeat); |
55 | } |
56 | |
57 | } // namespace duckdb |
58 | |