1//===----------------------------------------------------------------------===//
2// DuckDB
3//
4// duckdb/function/copy_function.hpp
5//
6//
7//===----------------------------------------------------------------------===//
8
9#pragma once
10
11#include "duckdb/function/function.hpp"
12#include "duckdb/function/table_function.hpp"
13#include "duckdb/parser/parsed_data/copy_info.hpp"
14#include "duckdb/parser/statement/copy_statement.hpp"
15
16namespace duckdb {
17
18class Binder;
19struct BoundStatement;
20class ColumnDataCollection;
21class ExecutionContext;
22
23struct LocalFunctionData {
24 virtual ~LocalFunctionData() {
25 }
26
27 template <class TARGET>
28 TARGET &Cast() {
29 D_ASSERT(dynamic_cast<TARGET *>(this));
30 return reinterpret_cast<TARGET &>(*this);
31 }
32 template <class TARGET>
33 const TARGET &Cast() const {
34 D_ASSERT(dynamic_cast<const TARGET *>(this));
35 return reinterpret_cast<const TARGET &>(*this);
36 }
37};
38
39struct GlobalFunctionData {
40 virtual ~GlobalFunctionData() {
41 }
42
43 template <class TARGET>
44 TARGET &Cast() {
45 D_ASSERT(dynamic_cast<TARGET *>(this));
46 return reinterpret_cast<TARGET &>(*this);
47 }
48 template <class TARGET>
49 const TARGET &Cast() const {
50 D_ASSERT(dynamic_cast<const TARGET *>(this));
51 return reinterpret_cast<const TARGET &>(*this);
52 }
53};
54
55struct PreparedBatchData {
56 virtual ~PreparedBatchData() {
57 }
58
59 template <class TARGET>
60 TARGET &Cast() {
61 D_ASSERT(dynamic_cast<TARGET *>(this));
62 return reinterpret_cast<TARGET &>(*this);
63 }
64 template <class TARGET>
65 const TARGET &Cast() const {
66 D_ASSERT(dynamic_cast<const TARGET *>(this));
67 return reinterpret_cast<const TARGET &>(*this);
68 }
69};
70
71enum class CopyFunctionExecutionMode { REGULAR_COPY_TO_FILE, PARALLEL_COPY_TO_FILE, BATCH_COPY_TO_FILE };
72
73typedef BoundStatement (*copy_to_plan_t)(Binder &binder, CopyStatement &stmt);
74typedef unique_ptr<FunctionData> (*copy_to_bind_t)(ClientContext &context, CopyInfo &info, vector<string> &names,
75 vector<LogicalType> &sql_types);
76typedef unique_ptr<LocalFunctionData> (*copy_to_initialize_local_t)(ExecutionContext &context, FunctionData &bind_data);
77typedef unique_ptr<GlobalFunctionData> (*copy_to_initialize_global_t)(ClientContext &context, FunctionData &bind_data,
78 const string &file_path);
79typedef void (*copy_to_sink_t)(ExecutionContext &context, FunctionData &bind_data, GlobalFunctionData &gstate,
80 LocalFunctionData &lstate, DataChunk &input);
81typedef void (*copy_to_combine_t)(ExecutionContext &context, FunctionData &bind_data, GlobalFunctionData &gstate,
82 LocalFunctionData &lstate);
83typedef void (*copy_to_finalize_t)(ClientContext &context, FunctionData &bind_data, GlobalFunctionData &gstate);
84
85typedef void (*copy_to_serialize_t)(FieldWriter &writer, const FunctionData &bind_data, const CopyFunction &function);
86
87typedef unique_ptr<FunctionData> (*copy_to_deserialize_t)(ClientContext &context, FieldReader &reader,
88 CopyFunction &function);
89
90typedef unique_ptr<FunctionData> (*copy_from_bind_t)(ClientContext &context, CopyInfo &info,
91 vector<string> &expected_names,
92 vector<LogicalType> &expected_types);
93typedef CopyFunctionExecutionMode (*copy_to_execution_mode_t)(bool preserve_insertion_order, bool supports_batch_index);
94
95typedef unique_ptr<PreparedBatchData> (*copy_prepare_batch_t)(ClientContext &context, FunctionData &bind_data,
96 GlobalFunctionData &gstate,
97 unique_ptr<ColumnDataCollection> collection);
98typedef void (*copy_flush_batch_t)(ClientContext &context, FunctionData &bind_data, GlobalFunctionData &gstate,
99 PreparedBatchData &batch);
100typedef idx_t (*copy_desired_batch_size_t)(ClientContext &context, FunctionData &bind_data);
101
102class CopyFunction : public Function {
103public:
104 explicit CopyFunction(string name)
105 : Function(name), plan(nullptr), copy_to_bind(nullptr), copy_to_initialize_local(nullptr),
106 copy_to_initialize_global(nullptr), copy_to_sink(nullptr), copy_to_combine(nullptr),
107 copy_to_finalize(nullptr), execution_mode(nullptr), prepare_batch(nullptr), flush_batch(nullptr),
108 desired_batch_size(nullptr), serialize(nullptr), deserialize(nullptr), copy_from_bind(nullptr) {
109 }
110
111 //! Plan rewrite copy function
112 copy_to_plan_t plan;
113
114 copy_to_bind_t copy_to_bind;
115 copy_to_initialize_local_t copy_to_initialize_local;
116 copy_to_initialize_global_t copy_to_initialize_global;
117 copy_to_sink_t copy_to_sink;
118 copy_to_combine_t copy_to_combine;
119 copy_to_finalize_t copy_to_finalize;
120 copy_to_execution_mode_t execution_mode;
121
122 copy_prepare_batch_t prepare_batch;
123 copy_flush_batch_t flush_batch;
124 copy_desired_batch_size_t desired_batch_size;
125
126 copy_to_serialize_t serialize;
127 copy_to_deserialize_t deserialize;
128
129 copy_from_bind_t copy_from_bind;
130 TableFunction copy_from_function;
131
132 string extension;
133};
134
135} // namespace duckdb
136