| 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 | |