/root/doris/be/src/udf/udf.h
Line | Count | Source (jump to first uncovered line) |
1 | | // Licensed to the Apache Software Foundation (ASF) under one |
2 | | // or more contributor license agreements. See the NOTICE file |
3 | | // distributed with this work for additional information |
4 | | // regarding copyright ownership. The ASF licenses this file |
5 | | // to you under the Apache License, Version 2.0 (the |
6 | | // "License"); you may not use this file except in compliance |
7 | | // with the License. You may obtain a copy of the License at |
8 | | // |
9 | | // http://www.apache.org/licenses/LICENSE-2.0 |
10 | | // |
11 | | // Unless required by applicable law or agreed to in writing, |
12 | | // software distributed under the License is distributed on an |
13 | | // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
14 | | // KIND, either express or implied. See the License for the |
15 | | // specific language governing permissions and limitations |
16 | | // under the License. |
17 | | // This file is copied from |
18 | | // https://github.com/apache/impala/blob/branch-2.9.0/be/src/udf/udf.h |
19 | | // and modified by Doris |
20 | | |
21 | | #pragma once |
22 | | |
23 | | #include <cstdint> |
24 | | #include <memory> |
25 | | #include <string> |
26 | | #include <vector> |
27 | | |
28 | | #include "runtime/types.h" |
29 | | #include "util/runtime_profile.h" |
30 | | #include "vec/common/arena.h" |
31 | | |
32 | | namespace doris { |
33 | | |
34 | | struct ColumnPtrWrapper; |
35 | | struct StringRef; |
36 | | class RuntimeState; |
37 | | |
38 | | // The FunctionContext is passed to every UDF/UDA and is the interface for the UDF to the |
39 | | // rest of the system. It contains APIs to examine the system state, report errors |
40 | | // and manage memory. |
41 | | class FunctionContext { |
42 | | public: |
43 | | enum FunctionStateScope { |
44 | | /// Indicates that the function state for this FunctionContext's UDF is shared across |
45 | | /// the plan fragment (a query is divided into multiple plan fragments, each of which |
46 | | /// is responsible for a part of the query execution). Within the plan fragment, there |
47 | | /// may be multiple instances of the UDF executing concurrently with multiple |
48 | | /// FunctionContexts sharing this state, meaning that the state must be |
49 | | /// thread-safe. The Prepare() function for the UDF may be called with this scope |
50 | | /// concurrently on a single host if the UDF will be evaluated in multiple plan |
51 | | /// fragments on that host. In general, read-only state that doesn't need to be |
52 | | /// recomputed for every UDF call should be fragment-local. |
53 | | /// TODO: not yet implemented |
54 | | FRAGMENT_LOCAL, |
55 | | |
56 | | /// Indicates that the function state is local to the execution thread. This state |
57 | | /// does not need to be thread-safe. However, this state will be initialized (via the |
58 | | /// Prepare() function) once for every execution thread, so fragment-local state |
59 | | /// should be used when possible for better performance. In general, inexpensive |
60 | | /// shared state that is written to by the UDF (e.g. scratch space) should be |
61 | | /// thread-local. |
62 | | THREAD_LOCAL, |
63 | | }; |
64 | | |
65 | | static std::unique_ptr<doris::FunctionContext> create_context( |
66 | | RuntimeState* state, const doris::TypeDescriptor& return_type, |
67 | | const std::vector<doris::TypeDescriptor>& arg_types); |
68 | | |
69 | | /// Returns a new FunctionContext with the same constant args, fragment-local state, and |
70 | | /// debug flag as this FunctionContext. The caller is responsible for calling delete on |
71 | | /// it. |
72 | | std::unique_ptr<doris::FunctionContext> clone(); |
73 | | |
74 | | void set_constant_cols(const std::vector<std::shared_ptr<doris::ColumnPtrWrapper>>& cols); |
75 | | |
76 | 33 | RuntimeState* state() { return _state; } |
77 | | |
78 | 6.57k | bool check_overflow_for_decimal() const { return _check_overflow_for_decimal; } |
79 | | |
80 | 47 | bool set_check_overflow_for_decimal(bool check_overflow_for_decimal) { |
81 | 47 | return _check_overflow_for_decimal = check_overflow_for_decimal; |
82 | 47 | } |
83 | | |
84 | 0 | void set_string_as_jsonb_string(bool string_as_jsonb_string) { |
85 | 0 | _string_as_jsonb_string = string_as_jsonb_string; |
86 | 0 | } |
87 | | |
88 | 0 | void set_jsonb_string_as_string(bool jsonb_string_as_string) { |
89 | 0 | _jsonb_string_as_string = jsonb_string_as_string; |
90 | 0 | } |
91 | | |
92 | 0 | void set_udf_execute_timer(RuntimeProfile::Counter* udf_execute_timer) { |
93 | 0 | _udf_execute_timer = udf_execute_timer; |
94 | 0 | } |
95 | | |
96 | 0 | RuntimeProfile::Counter* get_udf_execute_timer() { return _udf_execute_timer; } |
97 | | |
98 | | // Cast flag, when enable string_as_jsonb_string, string casting to jsonb will not parse string |
99 | | // instead just insert a string literal |
100 | 18 | bool string_as_jsonb_string() const { return _string_as_jsonb_string; } |
101 | | |
102 | | // Cast flag, when enable jsonb_string_as_string, jsonb string casting to string will not parse string |
103 | | // instead just insert a string literal |
104 | 112 | bool jsonb_string_as_string() const { return _jsonb_string_as_string; } |
105 | | |
106 | | // Sets an error for this UDF. If this is called, this will trigger the |
107 | | // query to fail. |
108 | | // Note: when you set error for the UDFs used in Data Load, you should |
109 | | // ensure the function return value is null. |
110 | | void set_error(const char* error_msg); |
111 | | |
112 | | // Adds a warning that is returned to the user. This can include things like |
113 | | // overflow or other recoverable error conditions. |
114 | | // Warnings are capped at a maximum number. Returns true if the warning was |
115 | | // added and false if it was ignored due to the cap. |
116 | | bool add_warning(const char* warning_msg); |
117 | | |
118 | | /// Methods for maintaining state across UDF/UDA function calls. SetFunctionState() can |
119 | | /// be used to store a pointer that can then be retrieved via GetFunctionState(). If |
120 | | /// GetFunctionState() is called when no pointer is set, it will return |
121 | | /// nullptr. SetFunctionState() does not take ownership of 'ptr'; it is up to the UDF/UDA |
122 | | /// to clean up any function state if necessary. |
123 | | void set_function_state(FunctionStateScope scope, std::shared_ptr<void> ptr); |
124 | | |
125 | | void* get_function_state(FunctionStateScope scope) const; |
126 | | |
127 | | // Returns the return type information of this function. For UDAs, this is the final |
128 | | // return type of the UDA (e.g., the type returned by the finalize function). |
129 | | const doris::TypeDescriptor& get_return_type() const; |
130 | | |
131 | | // Returns the number of arguments to this function (not including the FunctionContext* |
132 | | // argument). |
133 | | int get_num_args() const; |
134 | | |
135 | | // Returns the type information for the arg_idx-th argument (0-indexed, not including |
136 | | // the FunctionContext* argument). Returns nullptr if arg_idx is invalid. |
137 | | const doris::TypeDescriptor* get_arg_type(int arg_idx) const; |
138 | | |
139 | | // Returns true if the arg_idx-th input argument (0 indexed, not including the |
140 | | // FunctionContext* argument) is a constant (e.g. 5, "string", 1 + 1). |
141 | | bool is_col_constant(int arg_idx) const; |
142 | | |
143 | | // Returns a pointer to the value of the arg_idx-th input argument (0 indexed, not |
144 | | // including the FunctionContext* argument). Returns nullptr if the argument is not |
145 | | // constant. This function can be used to obtain user-specified constants in a UDF's |
146 | | // Init() or Close() functions. |
147 | | doris::ColumnPtrWrapper* get_constant_col(int arg_idx) const; |
148 | | |
149 | | // Creates a StringRef, which memory is available when this function context is used next time |
150 | | StringRef create_temp_string_val(int64_t len); |
151 | | |
152 | 11.0k | ~FunctionContext() = default; |
153 | | |
154 | 0 | vectorized::Arena& get_arena() { return arena; } |
155 | | |
156 | | private: |
157 | 11.0k | FunctionContext() = default; |
158 | | |
159 | | // Disable copy ctor and assignment operator |
160 | | FunctionContext(const FunctionContext& other); |
161 | | |
162 | | FunctionContext& operator=(const FunctionContext& other); |
163 | | |
164 | | // We use the query's runtime state to report errors and warnings. nullptr for test |
165 | | // contexts. |
166 | | RuntimeState* _state = nullptr; |
167 | | |
168 | | // Empty if there's no error |
169 | | std::string _error_msg; |
170 | | |
171 | | // The number of warnings reported. |
172 | | int64_t _num_warnings; |
173 | | |
174 | | /// The function state accessed via FunctionContext::Get/SetFunctionState() |
175 | | std::shared_ptr<void> _thread_local_fn_state; |
176 | | std::shared_ptr<void> _fragment_local_fn_state; |
177 | | |
178 | | // Type descriptor for the return type of the function. |
179 | | doris::TypeDescriptor _return_type; |
180 | | |
181 | | // Type descriptors for each argument of the function. |
182 | | std::vector<doris::TypeDescriptor> _arg_types; |
183 | | |
184 | | std::vector<std::shared_ptr<doris::ColumnPtrWrapper>> _constant_cols; |
185 | | |
186 | | //udf execute timer |
187 | | RuntimeProfile::Counter* _udf_execute_timer = nullptr; |
188 | | bool _check_overflow_for_decimal = false; |
189 | | |
190 | | bool _string_as_jsonb_string = false; |
191 | | bool _jsonb_string_as_string = false; |
192 | | |
193 | | std::string _string_result; |
194 | | |
195 | | vectorized::Arena arena; |
196 | | }; |
197 | | |
198 | | using doris::FunctionContext; |
199 | | } // namespace doris |