1 | //===----------------------------------------------------------------------===// |
2 | // DuckDB |
3 | // |
4 | // duckdb/main/query_result.hpp |
5 | // |
6 | // |
7 | //===----------------------------------------------------------------------===// |
8 | |
9 | #pragma once |
10 | |
11 | #include "duckdb/common/enums/statement_type.hpp" |
12 | #include "duckdb/common/types/data_chunk.hpp" |
13 | #include "duckdb/common/winapi.hpp" |
14 | #include "duckdb/common/preserved_error.hpp" |
15 | #include "duckdb/common/arrow/arrow_options.hpp" |
16 | |
17 | namespace duckdb { |
18 | struct BoxRendererConfig; |
19 | |
20 | enum class QueryResultType : uint8_t { MATERIALIZED_RESULT, STREAM_RESULT, PENDING_RESULT }; |
21 | |
22 | //! A set of properties from the client context that can be used to interpret the query result |
23 | struct ClientProperties { |
24 | ClientProperties(string time_zone_p, ArrowOffsetSize arrow_offset_size_p) |
25 | : time_zone(std::move(time_zone_p)), arrow_offset_size(arrow_offset_size_p) { |
26 | } |
27 | string time_zone; |
28 | ArrowOffsetSize arrow_offset_size; |
29 | }; |
30 | |
31 | class BaseQueryResult { |
32 | public: |
33 | //! Creates a successful query result with the specified names and types |
34 | DUCKDB_API BaseQueryResult(QueryResultType type, StatementType statement_type, StatementProperties properties, |
35 | vector<LogicalType> types, vector<string> names); |
36 | //! Creates an unsuccessful query result with error condition |
37 | DUCKDB_API BaseQueryResult(QueryResultType type, PreservedError error); |
38 | DUCKDB_API virtual ~BaseQueryResult(); |
39 | |
40 | //! The type of the result (MATERIALIZED or STREAMING) |
41 | QueryResultType type; |
42 | //! The type of the statement that created this result |
43 | StatementType statement_type; |
44 | //! Properties of the statement |
45 | StatementProperties properties; |
46 | //! The SQL types of the result |
47 | vector<LogicalType> types; |
48 | //! The names of the result |
49 | vector<string> names; |
50 | |
51 | public: |
52 | DUCKDB_API void ThrowError(const string &prepended_message = "" ) const; |
53 | DUCKDB_API void SetError(PreservedError error); |
54 | DUCKDB_API bool HasError() const; |
55 | DUCKDB_API const ExceptionType &GetErrorType() const; |
56 | DUCKDB_API const std::string &GetError(); |
57 | DUCKDB_API PreservedError &GetErrorObject(); |
58 | DUCKDB_API idx_t ColumnCount(); |
59 | |
60 | protected: |
61 | //! Whether or not execution was successful |
62 | bool success; |
63 | //! The error (in case execution was not successful) |
64 | PreservedError error; |
65 | }; |
66 | struct CurrentChunk { |
67 | //! The current data chunk |
68 | unique_ptr<DataChunk> data_chunk; |
69 | //! The current position in the data chunk |
70 | idx_t position; |
71 | //! If we have a current chunk we must scan for result production |
72 | bool Valid(); |
73 | //! The remaining size of the current chunk |
74 | idx_t RemainingSize(); |
75 | }; |
76 | //! The QueryResult object holds the result of a query. It can either be a MaterializedQueryResult, in which case the |
77 | //! result contains the entire result set, or a StreamQueryResult in which case the Fetch method can be called to |
78 | //! incrementally fetch data from the database. |
79 | class QueryResult : public BaseQueryResult { |
80 | public: |
81 | //! Creates a successful query result with the specified names and types |
82 | DUCKDB_API QueryResult(QueryResultType type, StatementType statement_type, StatementProperties properties, |
83 | vector<LogicalType> types, vector<string> names, ClientProperties client_properties); |
84 | //! Creates an unsuccessful query result with error condition |
85 | DUCKDB_API QueryResult(QueryResultType type, PreservedError error); |
86 | DUCKDB_API virtual ~QueryResult() override; |
87 | |
88 | //! Properties from the client context |
89 | ClientProperties client_properties; |
90 | //! The next result (if any) |
91 | unique_ptr<QueryResult> next; |
92 | //! In case we are converting the result from Native DuckDB to a different library (e.g., Arrow, Polars) |
93 | //! We might be producing chunks of a pre-determined size. |
94 | //! To comply, we use the following variable to store the current chunk, and it's position. |
95 | CurrentChunk current_chunk; |
96 | |
97 | public: |
98 | template <class TARGET> |
99 | TARGET &Cast() { |
100 | if (type != TARGET::TYPE) { |
101 | throw InternalException("Failed to cast query result to type - query result type mismatch" ); |
102 | } |
103 | return reinterpret_cast<TARGET &>(*this); |
104 | } |
105 | |
106 | template <class TARGET> |
107 | const TARGET &Cast() const { |
108 | if (type != TARGET::TYPE) { |
109 | throw InternalException("Failed to cast query result to type - query result type mismatch" ); |
110 | } |
111 | return reinterpret_cast<const TARGET &>(*this); |
112 | } |
113 | |
114 | public: |
115 | //! Returns the name of the column for the given index |
116 | DUCKDB_API const string &ColumnName(idx_t index) const; |
117 | //! Fetches a DataChunk of normalized (flat) vectors from the query result. |
118 | //! Returns nullptr if there are no more results to fetch. |
119 | DUCKDB_API virtual unique_ptr<DataChunk> Fetch(); |
120 | //! Fetches a DataChunk from the query result. The vectors are not normalized and hence any vector types can be |
121 | //! returned. |
122 | DUCKDB_API virtual unique_ptr<DataChunk> FetchRaw() = 0; |
123 | //! Converts the QueryResult to a string |
124 | DUCKDB_API virtual string ToString() = 0; |
125 | //! Converts the QueryResult to a box-rendered string |
126 | DUCKDB_API virtual string ToBox(ClientContext &context, const BoxRendererConfig &config); |
127 | //! Prints the QueryResult to the console |
128 | DUCKDB_API void Print(); |
129 | //! Returns true if the two results are identical; false otherwise. Note that this method is destructive; it calls |
130 | //! Fetch() until both results are exhausted. The data in the results will be lost. |
131 | DUCKDB_API bool Equals(QueryResult &other); |
132 | |
133 | bool TryFetch(unique_ptr<DataChunk> &result, PreservedError &error) { |
134 | try { |
135 | result = Fetch(); |
136 | return success; |
137 | } catch (const Exception &ex) { |
138 | error = PreservedError(ex); |
139 | return false; |
140 | } catch (std::exception &ex) { |
141 | error = PreservedError(ex); |
142 | return false; |
143 | } catch (...) { |
144 | error = PreservedError("Unknown error in Fetch" ); |
145 | return false; |
146 | } |
147 | } |
148 | |
149 | static ArrowOptions GetArrowOptions(QueryResult &query_result); |
150 | static string GetConfigTimezone(QueryResult &query_result); |
151 | |
152 | private: |
153 | class QueryResultIterator; |
154 | class QueryResultRow { |
155 | public: |
156 | explicit QueryResultRow(QueryResultIterator &iterator_p, idx_t row_idx) : iterator(iterator_p), row(0) { |
157 | } |
158 | |
159 | QueryResultIterator &iterator; |
160 | idx_t row; |
161 | |
162 | template <class T> |
163 | T GetValue(idx_t col_idx) const { |
164 | return iterator.chunk->GetValue(col_idx, index: row).GetValue<T>(); |
165 | } |
166 | }; |
167 | //! The row-based query result iterator. Invoking the |
168 | class QueryResultIterator { |
169 | public: |
170 | explicit QueryResultIterator(optional_ptr<QueryResult> result_p) |
171 | : current_row(*this, 0), result(result_p), base_row(0) { |
172 | if (result) { |
173 | chunk = shared_ptr<DataChunk>(result->Fetch().release()); |
174 | if (!chunk) { |
175 | result = nullptr; |
176 | } |
177 | } |
178 | } |
179 | |
180 | QueryResultRow current_row; |
181 | shared_ptr<DataChunk> chunk; |
182 | optional_ptr<QueryResult> result; |
183 | idx_t base_row; |
184 | |
185 | public: |
186 | void Next() { |
187 | if (!chunk) { |
188 | return; |
189 | } |
190 | current_row.row++; |
191 | if (current_row.row >= chunk->size()) { |
192 | base_row += chunk->size(); |
193 | chunk = shared_ptr<DataChunk>(result->Fetch().release()); |
194 | current_row.row = 0; |
195 | if (!chunk || chunk->size() == 0) { |
196 | // exhausted all rows |
197 | base_row = 0; |
198 | result = nullptr; |
199 | chunk.reset(); |
200 | } |
201 | } |
202 | } |
203 | |
204 | QueryResultIterator &operator++() { |
205 | Next(); |
206 | return *this; |
207 | } |
208 | bool operator!=(const QueryResultIterator &other) const { |
209 | return result != other.result || base_row != other.base_row || current_row.row != other.current_row.row; |
210 | } |
211 | const QueryResultRow &operator*() const { |
212 | return current_row; |
213 | } |
214 | }; |
215 | |
216 | public: |
217 | QueryResultIterator begin() { |
218 | return QueryResultIterator(this); |
219 | } |
220 | QueryResultIterator end() { |
221 | return QueryResultIterator(nullptr); |
222 | } |
223 | |
224 | protected: |
225 | DUCKDB_API string (); |
226 | |
227 | private: |
228 | QueryResult(const QueryResult &) = delete; |
229 | }; |
230 | |
231 | } // namespace duckdb |
232 | |