| 1 | #include "duckdb/main/materialized_query_result.hpp" |
| 2 | #include "duckdb/common/to_string.hpp" |
| 3 | #include "duckdb/main/client_context.hpp" |
| 4 | #include "duckdb/common/box_renderer.hpp" |
| 5 | |
| 6 | namespace duckdb { |
| 7 | |
| 8 | MaterializedQueryResult::MaterializedQueryResult(StatementType statement_type, StatementProperties properties, |
| 9 | vector<string> names_p, unique_ptr<ColumnDataCollection> collection_p, |
| 10 | ClientProperties client_properties) |
| 11 | : QueryResult(QueryResultType::MATERIALIZED_RESULT, statement_type, std::move(properties), collection_p->Types(), |
| 12 | std::move(names_p), std::move(client_properties)), |
| 13 | collection(std::move(collection_p)), scan_initialized(false) { |
| 14 | } |
| 15 | |
| 16 | MaterializedQueryResult::MaterializedQueryResult(PreservedError error) |
| 17 | : QueryResult(QueryResultType::MATERIALIZED_RESULT, std::move(error)), scan_initialized(false) { |
| 18 | } |
| 19 | |
| 20 | string MaterializedQueryResult::ToString() { |
| 21 | string result; |
| 22 | if (success) { |
| 23 | result = HeaderToString(); |
| 24 | result += "[ Rows: " + to_string(val: collection->Count()) + "]\n" ; |
| 25 | auto &coll = Collection(); |
| 26 | for (auto &row : coll.Rows()) { |
| 27 | for (idx_t col_idx = 0; col_idx < coll.ColumnCount(); col_idx++) { |
| 28 | if (col_idx > 0) { |
| 29 | result += "\t" ; |
| 30 | } |
| 31 | auto val = row.GetValue(column_index: col_idx); |
| 32 | result += val.IsNull() ? "NULL" : StringUtil::Replace(source: val.ToString(), from: string("\0" , 1), to: "\\0" ); |
| 33 | } |
| 34 | result += "\n" ; |
| 35 | } |
| 36 | result += "\n" ; |
| 37 | } else { |
| 38 | result = GetError() + "\n" ; |
| 39 | } |
| 40 | return result; |
| 41 | } |
| 42 | |
| 43 | string MaterializedQueryResult::ToBox(ClientContext &context, const BoxRendererConfig &config) { |
| 44 | if (!success) { |
| 45 | return GetError() + "\n" ; |
| 46 | } |
| 47 | if (!collection) { |
| 48 | return "Internal error - result was successful but there was no collection" ; |
| 49 | } |
| 50 | BoxRenderer renderer(config); |
| 51 | return renderer.ToString(context, names, op: Collection()); |
| 52 | } |
| 53 | |
| 54 | Value MaterializedQueryResult::GetValue(idx_t column, idx_t index) { |
| 55 | if (!row_collection) { |
| 56 | row_collection = make_uniq<ColumnDataRowCollection>(args: collection->GetRows()); |
| 57 | } |
| 58 | return row_collection->GetValue(column, index); |
| 59 | } |
| 60 | |
| 61 | idx_t MaterializedQueryResult::RowCount() const { |
| 62 | return collection ? collection->Count() : 0; |
| 63 | } |
| 64 | |
| 65 | ColumnDataCollection &MaterializedQueryResult::Collection() { |
| 66 | if (HasError()) { |
| 67 | throw InvalidInputException("Attempting to get collection from an unsuccessful query result\n: Error %s" , |
| 68 | GetError()); |
| 69 | } |
| 70 | if (!collection) { |
| 71 | throw InternalException("Missing collection from materialized query result" ); |
| 72 | } |
| 73 | return *collection; |
| 74 | } |
| 75 | |
| 76 | unique_ptr<DataChunk> MaterializedQueryResult::Fetch() { |
| 77 | return FetchRaw(); |
| 78 | } |
| 79 | |
| 80 | unique_ptr<DataChunk> MaterializedQueryResult::FetchRaw() { |
| 81 | if (HasError()) { |
| 82 | throw InvalidInputException("Attempting to fetch from an unsuccessful query result\nError: %s" , GetError()); |
| 83 | } |
| 84 | auto result = make_uniq<DataChunk>(); |
| 85 | collection->InitializeScanChunk(chunk&: *result); |
| 86 | if (!scan_initialized) { |
| 87 | // we disallow zero copy so the chunk is independently usable even after the result is destroyed |
| 88 | collection->InitializeScan(state&: scan_state, properties: ColumnDataScanProperties::DISALLOW_ZERO_COPY); |
| 89 | scan_initialized = true; |
| 90 | } |
| 91 | collection->Scan(state&: scan_state, result&: *result); |
| 92 | if (result->size() == 0) { |
| 93 | return nullptr; |
| 94 | } |
| 95 | return result; |
| 96 | } |
| 97 | |
| 98 | } // namespace duckdb |
| 99 | |