diff --git a/src/cut.cpp b/src/cut.cpp index d330169..a0c4595 100644 --- a/src/cut.cpp +++ b/src/cut.cpp @@ -50,11 +50,16 @@ int main( return 1; } - auto &query_plan = overall_query_plan->get_plans().back(); - query_plan.select.emplace( - query_plan.select->get_tablenames(), + auto &select = overall_query_plan->get_plans().back(); + if (!select.select.has_value()) { + std::cerr << "No SELECT fragment in final query plan.\n"; + return 1; + } + + select.select.emplace( + select.select->get_tablenames(), options.get_fields(), - query_plan.select->get_alias() + select.select->get_alias() ); return static_cast(dump_or_eval_query_plan(*overall_query_plan)); diff --git a/src/eval.cpp b/src/eval.cpp index f26eaae..c5ad3c4 100644 --- a/src/eval.cpp +++ b/src/eval.cpp @@ -28,8 +28,8 @@ class EvalOptions final : public Options { bool parse( const int argc, - const char *argv[] - ) override { // NOLINT(*-avoid-c-arrays) + const char *argv[] // NOLINT(*-avoid-c-arrays) + ) override { if (bool const parent_result = Options::parse(argc, argv); !parent_result) { return parent_result; } @@ -137,7 +137,7 @@ int main( } const auto [query_str, query_params] = *query; std::cout << query_str << '\n'; - for (const auto &[column, value] : query_params) { + for (const auto &[column, value]: query_params) { std::cout << "-- Colum " << column << ": " << value << '\n'; } return static_cast(ExitStatus::SUCCESS); diff --git a/src/grep.cpp b/src/grep.cpp index c91c58d..83054f7 100644 --- a/src/grep.cpp +++ b/src/grep.cpp @@ -28,8 +28,8 @@ class GrepOptions final : public Options { bool parse( const int argc, - const char *argv[] - ) override { // NOLINT(*-avoid-c-arrays) + const char *argv[] // NOLINT(*-avoid-c-arrays) + ) override { if (bool const parent_result = Options::parse(argc, argv); !parent_result) { return parent_result; } diff --git a/src/head.cpp b/src/head.cpp index a94273e..1d5e44c 100644 --- a/src/head.cpp +++ b/src/head.cpp @@ -23,8 +23,8 @@ class HeadOptions final : public Options { bool parse( const int argc, - const char *argv[] - ) override { // NOLINT(*-avoid-c-arrays) + const char *argv[] // NOLINT(*-avoid-c-arrays) + ) override { if (bool const parent_result = Options::parse(argc, argv); !parent_result) { return parent_result; } diff --git a/src/join.cpp b/src/join.cpp index e84b763..2fc871b 100644 --- a/src/join.cpp +++ b/src/join.cpp @@ -1,3 +1,4 @@ +#include #include #include @@ -61,7 +62,7 @@ class JoinOptions final : public Options { void parse_option_switches( std::vector &conditions ) const { - const auto num_args = std::min(left_.size(), std::min(right_.size(), predicate_.size())); + const auto num_args = std::min({left_.size(), right_.size(), predicate_.size()}); if (left_.size() != num_args || right_.size() != num_args || predicate_.size() != num_args) { std::cerr << "Number of left, right and predicate arguments must be the same. Got " << left_.size() << ", " << right_.size() << " and " << predicate_.size() << ". Using the first " << num_args << "\n"; @@ -84,8 +85,8 @@ class JoinOptions final : public Options { conditions.push_back( JoinCondition{ .left = positional_args_[i * 3], - .predicate = positional_args_[i * 3 + 1], - .right = positional_args_[i * 3 + 2] + .predicate = positional_args_[(i * 3) + 1], + .right = positional_args_[(i * 3) + 2] } ); } diff --git a/src/options.h b/src/options.h index 646e4bd..6c05cae 100644 --- a/src/options.h +++ b/src/options.h @@ -38,8 +38,8 @@ class Options { virtual bool parse( const int argc, - const char *argv[] - ) { // NOLINT(*-avoid-c-arrays) + const char *argv[] // NOLINT(*-avoid-c-arrays) + ) { namespace po = boost::program_options; try { diff --git a/src/query_evaluator.cpp b/src/query_evaluator.cpp index 281268a..7bdf4b2 100644 --- a/src/query_evaluator.cpp +++ b/src/query_evaluator.cpp @@ -3,11 +3,14 @@ #include #include #include +#include +#include +#include #include #include #include -#include +#include #include "arrow_result.h" #include "options.h" @@ -23,17 +26,103 @@ DuckDbException::DuckDbException( ) : std::runtime_error(msg) {} -template -T dd_check( - T result +namespace { +struct DatabaseCloser { + void operator()( + duckdb_database db + ) const { + if (db != nullptr) { + duckdb_close(&db); + } + } +}; + +struct ConnectionCloser { + void operator()( + duckdb_connection connection + ) const { + if (connection != nullptr) { + duckdb_disconnect(&connection); + } + } +}; + +struct PreparedStatementDestroyer { + void operator()( + duckdb_prepared_statement statement + ) const { + if (statement != nullptr) { + duckdb_destroy_prepare(&statement); + } + } +}; + +struct DataChunkDestroyer { + void operator()( + duckdb_data_chunk chunk + ) const { + if (chunk != nullptr) { + duckdb_destroy_data_chunk(&chunk); + } + } +}; + +struct ArrowOptionsDestroyer { + void operator()( + duckdb_arrow_options options + ) const { + if (options != nullptr) { + duckdb_destroy_arrow_options(&options); + } + } +}; + +struct DuckdbResult { + ~DuckdbResult() { + duckdb_destroy_result(&result_); + } + + [[nodiscard]] duckdb_result *get() { + return &result_; + } + duckdb_result result_{}; +}; + +using UniqueDatabase = std::unique_ptr, DatabaseCloser>; +using UniqueConnection = std::unique_ptr, ConnectionCloser>; +using UniquePreparedStatement = std::unique_ptr, + PreparedStatementDestroyer>; +using UniqueDataChunk = std::unique_ptr, DataChunkDestroyer>; +using UniqueArrowOptions = std::unique_ptr, ArrowOptionsDestroyer>; + + +[[noreturn]] void throw_duckdb_error( + const std::string &context, + const char *error_message +) { + if (error_message != nullptr) { + throw DuckDbException(context + " " + error_message); + } + throw DuckDbException(context); +} + +template +duckdb_state bind_checked( + const Binder &binder, + const duckdb_prepared_statement statement, + const std::string &context ) { - if (result->HasError()) { - throw DuckDbException("Error doing DuckDb action. " + result->GetErrorObject().Message()); + const auto state = binder(); + if (state == DuckDBError) { + throw_duckdb_error(context, duckdb_prepare_error(statement)); } - return result; + return state; } +} // namespace -static duckdb::Value infer_value_from_schema( +static void bind_inferred_value( + const duckdb_prepared_statement statement, + const idx_t param_index, const ColumnQueryParam ¶m, const std::unordered_map ¶m_types ) { @@ -45,97 +134,175 @@ static duckdb::Value infer_value_from_schema( if (col_type == "BIGINT" || col_type == "INTEGER" || col_type == "SMALLINT") { try { - return duckdb::Value::CreateValue(std::stoll(string_value)); + const auto integer_value = std::stoll(string_value); + bind_checked( + [&] { + return duckdb_bind_int64(statement, param_index, integer_value); + }, + statement, + "Could not bind inferred integer parameter."s + ); + return; } catch (const std::exception &e) { throw std::runtime_error( - "Could not convert paramter "s + string_value + " to integer for column '" + param.column + "': " + + "Could not convert parameter "s + string_value + " to integer for column '" + param.column + "': " + e.what() ); } } if (col_type == "DOUBLE" || col_type == "FLOAT") { try { - return duckdb::Value::CreateValue(std::stod(string_value)); + const auto double_value = std::stod(string_value); + bind_checked( + [&] { + return duckdb_bind_double(statement, param_index, double_value); + }, + statement, + "Could not bind inferred double parameter."s + ); + return; } catch (const std::exception &e) { throw std::runtime_error( - "Could not convert paramter "s + string_value + " to float for column '" + param.column + "': " + e. - what() + "Could not convert parameter "s + string_value + " to float for column '" + param.column + "': " + e + .what() ); } } if (col_type == "VARCHAR" || col_type == "TEXT") { - return duckdb::Value::CreateValue(param.value.get()); + const auto &string_param_value = param.value.get(); + bind_checked( + [&] { + return duckdb_bind_varchar(statement, param_index, string_param_value.c_str()); + }, + statement, + "Could not bind inferred text parameter."s + ); + return; } throw DuckDbException("Unable to infer type for param for column '" + param.column + "'."); } throw std::runtime_error("Could not find column '" + param.column + "' in schema."); } -static duckdb::vector convert_params_to_duckdb( +static void bind_params_to_statement( + const duckdb_prepared_statement statement, const std::vector &query_params, const std::unordered_map ¶m_types ) { - duckdb::vector duckdb_params; + using std::string_literals::operator ""s; - int i = 0; + idx_t param_index = 1; for (const auto ¶m: query_params) { - duckdb::Value value; switch (param.value.type()) { case ParamType::NUMERIC: - value = duckdb::Value::CreateValue(param.value.get()); + bind_checked( + [&] { + return duckdb_bind_int64(statement, param_index, param.value.get()); + }, + statement, + "Could not bind numeric parameter."s + ); break; - case ParamType::TEXT: - value = duckdb::Value::CreateValue(param.value.get()); + case ParamType::TEXT: { + const auto &string_value = param.value.get(); + bind_checked( + [&] { + return duckdb_bind_varchar(statement, param_index, string_value.c_str()); + }, + statement, + "Could not bind text parameter."s + ); break; + } case ParamType::UNKNOWN: - value = infer_value_from_schema(param, param_types); + bind_inferred_value(statement, param_index, param, param_types); break; default: - throw DuckDbException("Unable to convert param " + std::to_string(i) + " to DuckDb type."); + throw DuckDbException( + "Unable to convert param " + std::to_string(param_index - 1) + " to DuckDb type." + ); } - duckdb_params.push_back(value); - ++i; + ++param_index; } - - return duckdb_params; } -static std::shared_ptr duckdb_schema_to_arrow( - const std::unique_ptr &result -) { - ArrowSchema duck_arrow_schema{}; - duckdb::ArrowConverter::ToArrowSchema(&duck_arrow_schema, result->types, result->names, result->client_properties); +static auto duckdb_schema_to_arrow( + const duckdb_connection connection, + const duckdb_prepared_statement statement +) -> std::shared_ptr { + duckdb_arrow_options raw_arrow_options = nullptr; + duckdb_connection_get_arrow_options(connection, &raw_arrow_options); + if (raw_arrow_options == nullptr) { + throw DuckDbException("Could not fetch DuckDB Arrow options."); + } + const UniqueArrowOptions arrow_options(raw_arrow_options); + + const auto column_count = duckdb_prepared_statement_column_count(statement); + std::vector logical_types; + logical_types.reserve(column_count); + std::vector column_names; + column_names.reserve(column_count); + + for (idx_t idx = 0; idx < column_count; ++idx) { + logical_types.push_back(duckdb_prepared_statement_column_logical_type(statement, idx)); + column_names.push_back(duckdb_prepared_statement_column_name(statement, idx)); + } - return assign_or_raise(arrow::ImportSchema(&duck_arrow_schema)); + ArrowSchema arrow_schema{}; + duckdb_error_data error_data = duckdb_to_arrow_schema( + arrow_options.get(), + logical_types.data(), + column_names.data(), + column_count, + &arrow_schema + ); + + for (auto &logical_type: logical_types) { + duckdb_destroy_logical_type(&logical_type); + } + for (const auto &name: column_names) { + duckdb_free(const_cast(name)); // NOLINT(*-const-cast) + } + + if (duckdb_error_data_has_error(error_data)) { + const auto *message = duckdb_error_data_message(error_data); + duckdb_destroy_error_data(&error_data); + throw_duckdb_error("Could not convert schema to Arrow.", message); + } + duckdb_destroy_error_data(&error_data); + + return assign_or_raise(arrow::ImportSchema(&arrow_schema)); } -std::shared_ptr chunk_to_record_batch( - const std::unique_ptr &data_chunk, - std::shared_ptr arrow_schema, - const std::unique_ptr &result -) { +auto chunk_to_record_batch( + const duckdb_arrow_options arrow_options, + const duckdb_data_chunk chunk, + std::shared_ptr arrow_schema +) -> std::shared_ptr { ArrowArray arrow_array{}; - - // TODO: No need to recreate this multiple times - const std::unordered_map> - extension_type_cast; - duckdb::ArrowConverter::ToArrowArray(*data_chunk, &arrow_array, result->client_properties, extension_type_cast); + duckdb_error_data error_data = duckdb_data_chunk_to_arrow(arrow_options, chunk, &arrow_array); + if (duckdb_error_data_has_error(error_data)) { + const auto *message = duckdb_error_data_message(error_data); + duckdb_destroy_error_data(&error_data); + throw_duckdb_error("Could not convert chunk to Arrow.", message); + } + duckdb_destroy_error_data(&error_data); return assign_or_raise(arrow::ImportRecordBatch(&arrow_array, std::move(arrow_schema))); } static std::unordered_map get_schema( const OverallQueryPlan &query_plan, - duckdb::Connection &conn + duckdb_connection connection ) { using namespace std::string_literals; OverallQueryPlan base_query = query_plan; - if (base_query.get_plans().size() > 0) { + if (!base_query.get_plans().empty()) { auto &final_plan = base_query.get_plans().back(); final_plan.limit = std::nullopt; final_plan.order = std::nullopt; @@ -156,15 +323,27 @@ static std::unordered_map get_schema( std::unordered_map column_types; - const auto prepared_statement = dd_check(conn.Prepare(describe_query)); - const auto result = dd_check(prepared_statement->Execute()); + DuckdbResult result; + if (duckdb_query(connection, describe_query.c_str(), result.get()) == DuckDBError) { + const auto *error_message = duckdb_result_error(result.get()); + throw_duckdb_error("Error executing DESCRIBE query.", error_message); + } - for (auto data_chunk = result->Fetch(); data_chunk && data_chunk->size() > 0; data_chunk = result->Fetch()) { - for (duckdb::idx_t row = 0; row < data_chunk->size(); ++row) { - const auto column_name = data_chunk->GetValue(0, row).ToString(); - const auto column_type = data_chunk->GetValue(1, row).ToString(); - column_types[column_name] = column_type; + const auto row_count = duckdb_row_count(result.get()); + for (idx_t row = 0; row < row_count; ++row) { + const std::unique_ptr column_name( + duckdb_value_varchar(result.get(), 0, row), + duckdb_free + ); + const std::unique_ptr column_type( + duckdb_value_varchar(result.get(), 1, row), + duckdb_free + ); + + if (!column_name || !column_type) { + throw DuckDbException("Failed to read schema information from DuckDB."); } + column_types[column_name.get()] = column_type.get(); } return column_types; @@ -183,26 +362,56 @@ ExitStatus evaluate_query( return ExitStatus::QUERY_GENERATION_ERROR; } - duckdb::DuckDB db(nullptr); - duckdb::Connection con(db); + duckdb_database raw_db = nullptr; + if (duckdb_open(nullptr, &raw_db) == DuckDBError) { + throw DuckDbException("Failed to open DuckDB in-memory database."); + } + const UniqueDatabase database(raw_db); + + duckdb_connection raw_connection = nullptr; + if (duckdb_connect(database.get(), &raw_connection) == DuckDBError) { + throw DuckDbException("Failed to connect to DuckDB database."); + } + const UniqueConnection connection(raw_connection); auto [query_str, query_params] = *query; try { - const auto param_types = get_schema(query_plan, con); + const auto param_types = get_schema(query_plan, connection.get()); + + duckdb_prepared_statement raw_statement = nullptr; + const auto prepare_state = duckdb_prepare(connection.get(), query_str.c_str(), &raw_statement); + const UniquePreparedStatement prepared_statement(raw_statement); + if (prepare_state == DuckDBError) { + const auto *error_message = prepared_statement ? duckdb_prepare_error(prepared_statement.get()) : nullptr; + throw_duckdb_error("Failed to prepare statement.", error_message); + } - const auto prepared_statement = dd_check(con.Prepare(query_str)); - auto duckdb_params = convert_params_to_duckdb(query_params, param_types); - const auto result = dd_check(prepared_statement->Execute(duckdb_params, true)); - const auto arrow_schema = duckdb_schema_to_arrow(result); + bind_params_to_statement(prepared_statement.get(), query_params, param_types); + + DuckdbResult result; + if (duckdb_execute_prepared(prepared_statement.get(), result.get()) == DuckDBError) { + const auto *error_message = duckdb_result_error(result.get()); + throw_duckdb_error("Error executing prepared statement.", error_message); + } + + duckdb_arrow_options raw_arrow_options = duckdb_result_get_arrow_options(result.get()); + if (raw_arrow_options == nullptr) { + throw DuckDbException("DuckDB did not provide Arrow options for the result."); + } + const UniqueArrowOptions arrow_options(raw_arrow_options); + + const auto arrow_schema = duckdb_schema_to_arrow(connection.get(), prepared_statement.get()); const auto writer = writer_factory(arrow_schema); - while (true) { - auto data_chunk = result->Fetch(); - if (!data_chunk || data_chunk->size() == 0) { - break; + const auto chunk_count = duckdb_result_chunk_count(*result.get()); + for (idx_t chunk_index = 0; chunk_index < chunk_count; ++chunk_index) { + duckdb_data_chunk raw_chunk = duckdb_result_get_chunk(*result.get(), chunk_index); + const UniqueDataChunk chunk(raw_chunk); + if (!chunk || duckdb_data_chunk_get_size(chunk.get()) == 0) { + continue; } - const auto batch_result = chunk_to_record_batch(data_chunk, arrow_schema, result); + const auto batch_result = chunk_to_record_batch(arrow_options.get(), chunk.get(), arrow_schema); writer->write(batch_result); } writer->flush(); diff --git a/src/sort.cpp b/src/sort.cpp index 4b96f2f..5d81259 100644 --- a/src/sort.cpp +++ b/src/sort.cpp @@ -24,8 +24,8 @@ class SortOptions final : public Options { bool parse( const int argc, - const char *argv[] - ) override { // NOLINT(*-avoid-c-arrays) + const char *argv[] // NOLINT(*-avoid-c-arrays) + ) override { if (const bool parent_result = Options::parse(argc, argv); !parent_result) { return parent_result; } diff --git a/src/writer.cpp b/src/writer.cpp index da0db4f..8dc11bb 100644 --- a/src/writer.cpp +++ b/src/writer.cpp @@ -91,7 +91,7 @@ void ColumnarWriter::write( const auto cols = batch->columns(); std::int64_t num_rows = 0; - if (cols.size() == 0) { + if (cols.empty()) { throw std::runtime_error("No columns were provided"); } num_rows = cols[0]->length(); @@ -101,7 +101,7 @@ void ColumnarWriter::write( for (std::int64_t i = 0; i < num_rows; ++i) { for (std::size_t j = 0; j < cols.size(); ++j) { - const auto column = cols[j]; + const auto& column = cols[j]; const auto slice = column->Slice(i, 1); if (const auto status = arrow::PrettyPrint(*slice, print_options_, &stream); !status.ok()) { throw std::runtime_error("Error printing column: " + status.ToString()); @@ -124,8 +124,11 @@ void ColumnarWriter::write( void ColumnarWriter::flush() { for (const auto &row: rendered_rows_) { for (std::size_t i = 0; i < row.size(); ++i) { - const auto width = max_col_width_[i]; - const auto delimiter = i == 0 ? "" : " "; + const auto width = std::in_range(max_col_width_[i]) + ? static_cast(max_col_width_[i]) + : std::numeric_limits::max(); + + const auto *delimiter = i == 0 ? "" : " "; *stream_ << delimiter << std::left << std::setw(width); *stream_ << row[i]; }