1 | #include "duckdb/function/table/range.hpp" |
2 | #include "duckdb/common/algorithm.hpp" |
3 | #include "duckdb/planner/expression/bound_reference_expression.hpp" |
4 | #include "duckdb/planner/expression/bound_unnest_expression.hpp" |
5 | #include "duckdb/execution/operator/projection/physical_unnest.hpp" |
6 | |
7 | namespace duckdb { |
8 | |
9 | struct UnnestBindData : public FunctionData { |
10 | explicit UnnestBindData(LogicalType input_type_p) : input_type(std::move(input_type_p)) { |
11 | } |
12 | |
13 | LogicalType input_type; |
14 | |
15 | public: |
16 | unique_ptr<FunctionData> Copy() const override { |
17 | return make_uniq<UnnestBindData>(args: input_type); |
18 | } |
19 | |
20 | bool Equals(const FunctionData &other_p) const override { |
21 | auto &other = other_p.Cast<UnnestBindData>(); |
22 | return input_type == other.input_type; |
23 | } |
24 | }; |
25 | |
26 | struct UnnestGlobalState : public GlobalTableFunctionState { |
27 | UnnestGlobalState() { |
28 | } |
29 | |
30 | vector<unique_ptr<Expression>> select_list; |
31 | |
32 | idx_t MaxThreads() const override { |
33 | return GlobalTableFunctionState::MAX_THREADS; |
34 | } |
35 | }; |
36 | |
37 | struct UnnestLocalState : public LocalTableFunctionState { |
38 | UnnestLocalState() { |
39 | } |
40 | |
41 | unique_ptr<OperatorState> operator_state; |
42 | }; |
43 | |
44 | static unique_ptr<FunctionData> UnnestBind(ClientContext &context, TableFunctionBindInput &input, |
45 | vector<LogicalType> &return_types, vector<string> &names) { |
46 | if (input.input_table_types.size() != 1 || input.input_table_types[0].id() != LogicalTypeId::LIST) { |
47 | throw BinderException("UNNEST requires a single list as input" ); |
48 | } |
49 | return_types.push_back(x: ListType::GetChildType(type: input.input_table_types[0])); |
50 | names.push_back(x: input.input_table_names[0]); |
51 | return make_uniq<UnnestBindData>(args&: input.input_table_types[0]); |
52 | } |
53 | |
54 | static unique_ptr<LocalTableFunctionState> UnnestLocalInit(ExecutionContext &context, TableFunctionInitInput &input, |
55 | GlobalTableFunctionState *global_state) { |
56 | auto &gstate = global_state->Cast<UnnestGlobalState>(); |
57 | |
58 | auto result = make_uniq<UnnestLocalState>(); |
59 | result->operator_state = PhysicalUnnest::GetState(context, select_list: gstate.select_list); |
60 | return std::move(result); |
61 | } |
62 | |
63 | static unique_ptr<GlobalTableFunctionState> UnnestInit(ClientContext &context, TableFunctionInitInput &input) { |
64 | auto &bind_data = input.bind_data->Cast<UnnestBindData>(); |
65 | auto result = make_uniq<UnnestGlobalState>(); |
66 | auto ref = make_uniq<BoundReferenceExpression>(args: bind_data.input_type, args: 0); |
67 | auto bound_unnest = make_uniq<BoundUnnestExpression>(args: ListType::GetChildType(type: bind_data.input_type)); |
68 | bound_unnest->child = std::move(ref); |
69 | result->select_list.push_back(x: std::move(bound_unnest)); |
70 | return std::move(result); |
71 | } |
72 | |
73 | static OperatorResultType UnnestFunction(ExecutionContext &context, TableFunctionInput &data_p, DataChunk &input, |
74 | DataChunk &output) { |
75 | auto &state = data_p.global_state->Cast<UnnestGlobalState>(); |
76 | auto &lstate = data_p.local_state->Cast<UnnestLocalState>(); |
77 | return PhysicalUnnest::ExecuteInternal(context, input, chunk&: output, state&: *lstate.operator_state, select_list: state.select_list, include_input: false); |
78 | } |
79 | |
80 | void UnnestTableFunction::RegisterFunction(BuiltinFunctions &set) { |
81 | TableFunction unnest_function("unnest" , {LogicalTypeId::TABLE}, nullptr, UnnestBind, UnnestInit, UnnestLocalInit); |
82 | unnest_function.in_out_function = UnnestFunction; |
83 | set.AddFunction(function: unnest_function); |
84 | } |
85 | |
86 | } // namespace duckdb |
87 | |