1 | #include "duckdb/main/capi/capi_internal.hpp" |
2 | #include "duckdb/function/table_function.hpp" |
3 | #include "duckdb/parser/parsed_data/create_table_function_info.hpp" |
4 | #include "duckdb/catalog/catalog.hpp" |
5 | #include "duckdb/main/client_context.hpp" |
6 | #include "duckdb/storage/statistics/node_statistics.hpp" |
7 | |
8 | namespace duckdb { |
9 | |
10 | struct CTableFunctionInfo : public TableFunctionInfo { |
11 | ~CTableFunctionInfo() { |
12 | if (extra_info && delete_callback) { |
13 | delete_callback(extra_info); |
14 | } |
15 | extra_info = nullptr; |
16 | delete_callback = nullptr; |
17 | } |
18 | |
19 | duckdb_table_function_bind_t bind = nullptr; |
20 | duckdb_table_function_init_t init = nullptr; |
21 | duckdb_table_function_init_t local_init = nullptr; |
22 | duckdb_table_function_t function = nullptr; |
23 | void * = nullptr; |
24 | duckdb_delete_callback_t delete_callback = nullptr; |
25 | }; |
26 | |
27 | struct CTableBindData : public TableFunctionData { |
28 | CTableBindData(CTableFunctionInfo &info) : info(info) { |
29 | } |
30 | ~CTableBindData() { |
31 | if (bind_data && delete_callback) { |
32 | delete_callback(bind_data); |
33 | } |
34 | bind_data = nullptr; |
35 | delete_callback = nullptr; |
36 | } |
37 | |
38 | CTableFunctionInfo &info; |
39 | void *bind_data = nullptr; |
40 | duckdb_delete_callback_t delete_callback = nullptr; |
41 | unique_ptr<NodeStatistics> stats; |
42 | }; |
43 | |
44 | struct CTableInternalBindInfo { |
45 | CTableInternalBindInfo(ClientContext &context, TableFunctionBindInput &input, vector<LogicalType> &return_types, |
46 | vector<string> &names, CTableBindData &bind_data, CTableFunctionInfo &function_info) |
47 | : context(context), input(input), return_types(return_types), names(names), bind_data(bind_data), |
48 | function_info(function_info), success(true) { |
49 | } |
50 | |
51 | ClientContext &context; |
52 | TableFunctionBindInput &input; |
53 | vector<LogicalType> &return_types; |
54 | vector<string> &names; |
55 | CTableBindData &bind_data; |
56 | CTableFunctionInfo &function_info; |
57 | bool success; |
58 | string error; |
59 | }; |
60 | |
61 | struct CTableInitData { |
62 | ~CTableInitData() { |
63 | if (init_data && delete_callback) { |
64 | delete_callback(init_data); |
65 | } |
66 | init_data = nullptr; |
67 | delete_callback = nullptr; |
68 | } |
69 | |
70 | void *init_data = nullptr; |
71 | duckdb_delete_callback_t delete_callback = nullptr; |
72 | idx_t max_threads = 1; |
73 | }; |
74 | |
75 | struct CTableGlobalInitData : public GlobalTableFunctionState { |
76 | CTableInitData init_data; |
77 | |
78 | idx_t MaxThreads() const override { |
79 | return init_data.max_threads; |
80 | } |
81 | }; |
82 | |
83 | struct CTableLocalInitData : public LocalTableFunctionState { |
84 | CTableInitData init_data; |
85 | }; |
86 | |
87 | struct CTableInternalInitInfo { |
88 | CTableInternalInitInfo(const CTableBindData &bind_data, CTableInitData &init_data, |
89 | const vector<column_t> &column_ids, optional_ptr<TableFilterSet> filters) |
90 | : bind_data(bind_data), init_data(init_data), column_ids(column_ids), filters(filters), success(true) { |
91 | } |
92 | |
93 | const CTableBindData &bind_data; |
94 | CTableInitData &init_data; |
95 | const vector<column_t> &column_ids; |
96 | optional_ptr<TableFilterSet> filters; |
97 | bool success; |
98 | string error; |
99 | }; |
100 | |
101 | struct CTableInternalFunctionInfo { |
102 | CTableInternalFunctionInfo(const CTableBindData &bind_data, CTableInitData &init_data, CTableInitData &local_data) |
103 | : bind_data(bind_data), init_data(init_data), local_data(local_data), success(true) { |
104 | } |
105 | |
106 | const CTableBindData &bind_data; |
107 | CTableInitData &init_data; |
108 | CTableInitData &local_data; |
109 | bool success; |
110 | string error; |
111 | }; |
112 | |
113 | unique_ptr<FunctionData> CTableFunctionBind(ClientContext &context, TableFunctionBindInput &input, |
114 | vector<LogicalType> &return_types, vector<string> &names) { |
115 | auto &info = input.info->Cast<CTableFunctionInfo>(); |
116 | D_ASSERT(info.bind && info.function && info.init); |
117 | auto result = make_uniq<CTableBindData>(args&: info); |
118 | CTableInternalBindInfo bind_info(context, input, return_types, names, *result, info); |
119 | info.bind(&bind_info); |
120 | if (!bind_info.success) { |
121 | throw Exception(bind_info.error); |
122 | } |
123 | |
124 | return std::move(result); |
125 | } |
126 | |
127 | unique_ptr<GlobalTableFunctionState> CTableFunctionInit(ClientContext &context, TableFunctionInitInput &data_p) { |
128 | auto &bind_data = data_p.bind_data->Cast<CTableBindData>(); |
129 | auto result = make_uniq<CTableGlobalInitData>(); |
130 | |
131 | CTableInternalInitInfo init_info(bind_data, result->init_data, data_p.column_ids, data_p.filters); |
132 | bind_data.info.init(&init_info); |
133 | if (!init_info.success) { |
134 | throw Exception(init_info.error); |
135 | } |
136 | return std::move(result); |
137 | } |
138 | |
139 | unique_ptr<LocalTableFunctionState> CTableFunctionLocalInit(ExecutionContext &context, TableFunctionInitInput &data_p, |
140 | GlobalTableFunctionState *gstate) { |
141 | auto &bind_data = data_p.bind_data->Cast<CTableBindData>(); |
142 | auto result = make_uniq<CTableLocalInitData>(); |
143 | if (!bind_data.info.local_init) { |
144 | return std::move(result); |
145 | } |
146 | |
147 | CTableInternalInitInfo init_info(bind_data, result->init_data, data_p.column_ids, data_p.filters); |
148 | bind_data.info.local_init(&init_info); |
149 | if (!init_info.success) { |
150 | throw Exception(init_info.error); |
151 | } |
152 | return std::move(result); |
153 | } |
154 | |
155 | unique_ptr<NodeStatistics> CTableFunctionCardinality(ClientContext &context, const FunctionData *bind_data_p) { |
156 | auto &bind_data = bind_data_p->Cast<CTableBindData>(); |
157 | if (!bind_data.stats) { |
158 | return nullptr; |
159 | } |
160 | return make_uniq<NodeStatistics>(args&: *bind_data.stats); |
161 | } |
162 | |
163 | void CTableFunction(ClientContext &context, TableFunctionInput &data_p, DataChunk &output) { |
164 | auto &bind_data = data_p.bind_data->Cast<CTableBindData>(); |
165 | auto &global_data = (CTableGlobalInitData &)*data_p.global_state; |
166 | auto &local_data = (CTableLocalInitData &)*data_p.local_state; |
167 | CTableInternalFunctionInfo function_info(bind_data, global_data.init_data, local_data.init_data); |
168 | bind_data.info.function(&function_info, reinterpret_cast<duckdb_data_chunk>(&output)); |
169 | if (!function_info.success) { |
170 | throw Exception(function_info.error); |
171 | } |
172 | } |
173 | |
174 | } // namespace duckdb |
175 | |
176 | //===--------------------------------------------------------------------===// |
177 | // Table Function |
178 | //===--------------------------------------------------------------------===// |
179 | duckdb_table_function duckdb_create_table_function() { |
180 | auto function = new duckdb::TableFunction("" , {}, duckdb::CTableFunction, duckdb::CTableFunctionBind, |
181 | duckdb::CTableFunctionInit, duckdb::CTableFunctionLocalInit); |
182 | function->function_info = duckdb::make_shared<duckdb::CTableFunctionInfo>(); |
183 | function->cardinality = duckdb::CTableFunctionCardinality; |
184 | return function; |
185 | } |
186 | |
187 | void duckdb_destroy_table_function(duckdb_table_function *function) { |
188 | if (function && *function) { |
189 | auto tf = (duckdb::TableFunction *)*function; |
190 | delete tf; |
191 | *function = nullptr; |
192 | } |
193 | } |
194 | |
195 | void duckdb_table_function_set_name(duckdb_table_function function, const char *name) { |
196 | if (!function || !name) { |
197 | return; |
198 | } |
199 | auto tf = (duckdb::TableFunction *)function; |
200 | tf->name = name; |
201 | } |
202 | |
203 | void duckdb_table_function_add_parameter(duckdb_table_function function, duckdb_logical_type type) { |
204 | if (!function || !type) { |
205 | return; |
206 | } |
207 | auto tf = (duckdb::TableFunction *)function; |
208 | auto logical_type = (duckdb::LogicalType *)type; |
209 | tf->arguments.push_back(x: *logical_type); |
210 | } |
211 | |
212 | void duckdb_table_function_add_named_parameter(duckdb_table_function function, const char *name, |
213 | duckdb_logical_type type) { |
214 | if (!function || !type) { |
215 | return; |
216 | } |
217 | auto tf = (duckdb::TableFunction *)function; |
218 | auto logical_type = (duckdb::LogicalType *)type; |
219 | tf->named_parameters.insert(x: {name, *logical_type}); |
220 | } |
221 | |
222 | void (duckdb_table_function function, void *, |
223 | duckdb_delete_callback_t destroy) { |
224 | if (!function) { |
225 | return; |
226 | } |
227 | auto tf = (duckdb::TableFunction *)function; |
228 | auto info = (duckdb::CTableFunctionInfo *)tf->function_info.get(); |
229 | info->extra_info = extra_info; |
230 | info->delete_callback = destroy; |
231 | } |
232 | |
233 | void duckdb_table_function_set_bind(duckdb_table_function function, duckdb_table_function_bind_t bind) { |
234 | if (!function || !bind) { |
235 | return; |
236 | } |
237 | auto tf = (duckdb::TableFunction *)function; |
238 | auto info = (duckdb::CTableFunctionInfo *)tf->function_info.get(); |
239 | info->bind = bind; |
240 | } |
241 | |
242 | void duckdb_table_function_set_init(duckdb_table_function function, duckdb_table_function_init_t init) { |
243 | if (!function || !init) { |
244 | return; |
245 | } |
246 | auto tf = (duckdb::TableFunction *)function; |
247 | auto info = (duckdb::CTableFunctionInfo *)tf->function_info.get(); |
248 | info->init = init; |
249 | } |
250 | |
251 | void duckdb_table_function_set_local_init(duckdb_table_function function, duckdb_table_function_init_t init) { |
252 | if (!function || !init) { |
253 | return; |
254 | } |
255 | auto tf = (duckdb::TableFunction *)function; |
256 | auto info = (duckdb::CTableFunctionInfo *)tf->function_info.get(); |
257 | info->local_init = init; |
258 | } |
259 | |
260 | void duckdb_table_function_set_function(duckdb_table_function table_function, duckdb_table_function_t function) { |
261 | if (!table_function || !function) { |
262 | return; |
263 | } |
264 | auto tf = (duckdb::TableFunction *)table_function; |
265 | auto info = (duckdb::CTableFunctionInfo *)tf->function_info.get(); |
266 | info->function = function; |
267 | } |
268 | |
269 | void duckdb_table_function_supports_projection_pushdown(duckdb_table_function table_function, bool pushdown) { |
270 | if (!table_function) { |
271 | return; |
272 | } |
273 | auto tf = (duckdb::TableFunction *)table_function; |
274 | tf->projection_pushdown = pushdown; |
275 | } |
276 | |
277 | duckdb_state duckdb_register_table_function(duckdb_connection connection, duckdb_table_function function) { |
278 | if (!connection || !function) { |
279 | return DuckDBError; |
280 | } |
281 | auto con = (duckdb::Connection *)connection; |
282 | auto tf = (duckdb::TableFunction *)function; |
283 | auto info = (duckdb::CTableFunctionInfo *)tf->function_info.get(); |
284 | if (tf->name.empty() || !info->bind || !info->init || !info->function) { |
285 | return DuckDBError; |
286 | } |
287 | con->context->RunFunctionInTransaction(fun: [&]() { |
288 | auto &catalog = duckdb::Catalog::GetSystemCatalog(context&: *con->context); |
289 | duckdb::CreateTableFunctionInfo tf_info(*tf); |
290 | |
291 | // create the function in the catalog |
292 | catalog.CreateTableFunction(context&: *con->context, info&: tf_info); |
293 | }); |
294 | return DuckDBSuccess; |
295 | } |
296 | |
297 | //===--------------------------------------------------------------------===// |
298 | // Bind Interface |
299 | //===--------------------------------------------------------------------===// |
300 | void *(duckdb_bind_info info) { |
301 | if (!info) { |
302 | return nullptr; |
303 | } |
304 | auto bind_info = (duckdb::CTableInternalBindInfo *)info; |
305 | return bind_info->function_info.extra_info; |
306 | } |
307 | |
308 | void duckdb_bind_add_result_column(duckdb_bind_info info, const char *name, duckdb_logical_type type) { |
309 | if (!info || !name || !type) { |
310 | return; |
311 | } |
312 | auto bind_info = (duckdb::CTableInternalBindInfo *)info; |
313 | bind_info->names.push_back(x: name); |
314 | bind_info->return_types.push_back(x: *(reinterpret_cast<duckdb::LogicalType *>(type))); |
315 | } |
316 | |
317 | idx_t duckdb_bind_get_parameter_count(duckdb_bind_info info) { |
318 | if (!info) { |
319 | return 0; |
320 | } |
321 | auto bind_info = (duckdb::CTableInternalBindInfo *)info; |
322 | return bind_info->input.inputs.size(); |
323 | } |
324 | |
325 | duckdb_value duckdb_bind_get_parameter(duckdb_bind_info info, idx_t index) { |
326 | if (!info || index >= duckdb_bind_get_parameter_count(info)) { |
327 | return nullptr; |
328 | } |
329 | auto bind_info = (duckdb::CTableInternalBindInfo *)info; |
330 | return reinterpret_cast<duckdb_value>(new duckdb::Value(bind_info->input.inputs[index])); |
331 | } |
332 | |
333 | duckdb_value duckdb_bind_get_named_parameter(duckdb_bind_info info, const char *name) { |
334 | if (!info || !name) { |
335 | return nullptr; |
336 | } |
337 | auto bind_info = (duckdb::CTableInternalBindInfo *)info; |
338 | auto t = bind_info->input.named_parameters.find(x: name); |
339 | if (t == bind_info->input.named_parameters.end()) { |
340 | return nullptr; |
341 | } else { |
342 | return reinterpret_cast<duckdb_value>(new duckdb::Value(t->second)); |
343 | } |
344 | } |
345 | |
346 | void duckdb_bind_set_bind_data(duckdb_bind_info info, void *bind_data, duckdb_delete_callback_t destroy) { |
347 | if (!info) { |
348 | return; |
349 | } |
350 | auto bind_info = (duckdb::CTableInternalBindInfo *)info; |
351 | bind_info->bind_data.bind_data = bind_data; |
352 | bind_info->bind_data.delete_callback = destroy; |
353 | } |
354 | |
355 | void duckdb_bind_set_cardinality(duckdb_bind_info info, idx_t cardinality, bool is_exact) { |
356 | if (!info) { |
357 | return; |
358 | } |
359 | auto bind_info = (duckdb::CTableInternalBindInfo *)info; |
360 | if (is_exact) { |
361 | bind_info->bind_data.stats = duckdb::make_uniq<duckdb::NodeStatistics>(args&: cardinality); |
362 | } else { |
363 | bind_info->bind_data.stats = duckdb::make_uniq<duckdb::NodeStatistics>(args&: cardinality, args&: cardinality); |
364 | } |
365 | } |
366 | |
367 | void duckdb_bind_set_error(duckdb_bind_info info, const char *error) { |
368 | if (!info || !error) { |
369 | return; |
370 | } |
371 | auto function_info = (duckdb::CTableInternalBindInfo *)info; |
372 | function_info->error = error; |
373 | function_info->success = false; |
374 | } |
375 | |
376 | //===--------------------------------------------------------------------===// |
377 | // Init Interface |
378 | //===--------------------------------------------------------------------===// |
379 | void *(duckdb_init_info info) { |
380 | if (!info) { |
381 | return nullptr; |
382 | } |
383 | auto init_info = (duckdb::CTableInternalInitInfo *)info; |
384 | return init_info->bind_data.info.extra_info; |
385 | } |
386 | |
387 | void *duckdb_init_get_bind_data(duckdb_init_info info) { |
388 | if (!info) { |
389 | return nullptr; |
390 | } |
391 | auto init_info = (duckdb::CTableInternalInitInfo *)info; |
392 | return init_info->bind_data.bind_data; |
393 | } |
394 | |
395 | void duckdb_init_set_init_data(duckdb_init_info info, void *init_data, duckdb_delete_callback_t destroy) { |
396 | if (!info) { |
397 | return; |
398 | } |
399 | auto init_info = (duckdb::CTableInternalInitInfo *)info; |
400 | init_info->init_data.init_data = init_data; |
401 | init_info->init_data.delete_callback = destroy; |
402 | } |
403 | |
404 | void duckdb_init_set_error(duckdb_init_info info, const char *error) { |
405 | if (!info || !error) { |
406 | return; |
407 | } |
408 | auto function_info = (duckdb::CTableInternalInitInfo *)info; |
409 | function_info->error = error; |
410 | function_info->success = false; |
411 | } |
412 | |
413 | idx_t duckdb_init_get_column_count(duckdb_init_info info) { |
414 | if (!info) { |
415 | return 0; |
416 | } |
417 | auto function_info = (duckdb::CTableInternalInitInfo *)info; |
418 | return function_info->column_ids.size(); |
419 | } |
420 | |
421 | idx_t duckdb_init_get_column_index(duckdb_init_info info, idx_t column_index) { |
422 | if (!info) { |
423 | return 0; |
424 | } |
425 | auto function_info = (duckdb::CTableInternalInitInfo *)info; |
426 | if (column_index >= function_info->column_ids.size()) { |
427 | return 0; |
428 | } |
429 | return function_info->column_ids[column_index]; |
430 | } |
431 | |
432 | void duckdb_init_set_max_threads(duckdb_init_info info, idx_t max_threads) { |
433 | if (!info) { |
434 | return; |
435 | } |
436 | auto function_info = (duckdb::CTableInternalInitInfo *)info; |
437 | function_info->init_data.max_threads = max_threads; |
438 | } |
439 | |
440 | //===--------------------------------------------------------------------===// |
441 | // Function Interface |
442 | //===--------------------------------------------------------------------===// |
443 | void *(duckdb_function_info info) { |
444 | if (!info) { |
445 | return nullptr; |
446 | } |
447 | auto function_info = (duckdb::CTableInternalFunctionInfo *)info; |
448 | return function_info->bind_data.info.extra_info; |
449 | } |
450 | |
451 | void *duckdb_function_get_bind_data(duckdb_function_info info) { |
452 | if (!info) { |
453 | return nullptr; |
454 | } |
455 | auto function_info = (duckdb::CTableInternalFunctionInfo *)info; |
456 | return function_info->bind_data.bind_data; |
457 | } |
458 | |
459 | void *duckdb_function_get_init_data(duckdb_function_info info) { |
460 | if (!info) { |
461 | return nullptr; |
462 | } |
463 | auto function_info = (duckdb::CTableInternalFunctionInfo *)info; |
464 | return function_info->init_data.init_data; |
465 | } |
466 | |
467 | void *duckdb_function_get_local_init_data(duckdb_function_info info) { |
468 | if (!info) { |
469 | return nullptr; |
470 | } |
471 | auto function_info = (duckdb::CTableInternalFunctionInfo *)info; |
472 | return function_info->local_data.init_data; |
473 | } |
474 | |
475 | void duckdb_function_set_error(duckdb_function_info info, const char *error) { |
476 | if (!info || !error) { |
477 | return; |
478 | } |
479 | auto function_info = (duckdb::CTableInternalFunctionInfo *)info; |
480 | function_info->error = error; |
481 | function_info->success = false; |
482 | } |
483 | |