| 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 | |