1 | #include "duckdb/function/table/read_csv.hpp" |
2 | #include "duckdb/execution/operator/persistent/buffered_csv_reader.hpp" |
3 | |
4 | using namespace std; |
5 | |
6 | namespace duckdb { |
7 | |
8 | struct ReadCSVData : public TableFunctionData { |
9 | ReadCSVData() { |
10 | } |
11 | |
12 | CopyInfo info; |
13 | //! The CSV reader |
14 | unique_ptr<BufferedCSVReader> csv_reader; |
15 | }; |
16 | |
17 | static unique_ptr<FunctionData> read_csv_bind(ClientContext &context, vector<Value> inputs, |
18 | vector<SQLType> &return_types, vector<string> &names) { |
19 | for (auto &val : inputs[2].struct_value) { |
20 | names.push_back(val.first); |
21 | if (val.second.type != TypeId::VARCHAR) { |
22 | throw BinderException("read_csv requires a type specification as string" ); |
23 | } |
24 | return_types.push_back(TransformStringToSQLType(val.second.str_value.c_str())); |
25 | } |
26 | if (names.size() == 0) { |
27 | throw BinderException("read_csv requires at least a single column as input!" ); |
28 | } |
29 | auto result = make_unique<ReadCSVData>(); |
30 | |
31 | result->info.file_path = inputs[0].str_value; |
32 | result->info.header = false; |
33 | result->info.delimiter = inputs[1].str_value; |
34 | |
35 | result->csv_reader = make_unique<BufferedCSVReader>(context, result->info, return_types); |
36 | return move(result); |
37 | } |
38 | |
39 | static unique_ptr<FunctionData> read_csv_auto_bind(ClientContext &context, vector<Value> inputs, |
40 | vector<SQLType> &return_types, vector<string> &names) { |
41 | auto result = make_unique<ReadCSVData>(); |
42 | result->info.file_path = inputs[0].str_value; |
43 | result->info.auto_detect = true; |
44 | |
45 | result->csv_reader = make_unique<BufferedCSVReader>(context, result->info); |
46 | |
47 | // TODO: print detected dialect from result->csv_reader->info |
48 | return_types.assign(result->csv_reader->sql_types.begin(), result->csv_reader->sql_types.end()); |
49 | names.assign(result->csv_reader->col_names.begin(), result->csv_reader->col_names.end()); |
50 | |
51 | return move(result); |
52 | } |
53 | |
54 | static void read_csv_info(ClientContext &context, vector<Value> &input, DataChunk &output, FunctionData *dataptr) { |
55 | auto &data = ((ReadCSVData &)*dataptr); |
56 | data.csv_reader->ParseCSV(output); |
57 | } |
58 | |
59 | void ReadCSVTableFunction::RegisterFunction(BuiltinFunctions &set) { |
60 | set.AddFunction(TableFunction("read_csv" , {SQLType::VARCHAR, SQLType::VARCHAR, SQLType::STRUCT}, read_csv_bind, |
61 | read_csv_info, nullptr)); |
62 | set.AddFunction(TableFunction("read_csv_auto" , {SQLType::VARCHAR}, read_csv_auto_bind, read_csv_info, nullptr)); |
63 | } |
64 | |
65 | void BuiltinFunctions::RegisterReadFunctions() { |
66 | ReadCSVTableFunction::RegisterFunction(*this); |
67 | } |
68 | |
69 | } // namespace duckdb |
70 | |