be/src/format_v2/table/adbc_reader.cpp
Line | Count | Source |
1 | | // Licensed to the Apache Software Foundation (ASF) under one |
2 | | // or more contributor license agreements. See the NOTICE file |
3 | | // distributed with this work for additional information |
4 | | // regarding copyright ownership. The ASF licenses this file |
5 | | // to you under the Apache License, Version 2.0 (the |
6 | | // "License"); you may not use this file except in compliance |
7 | | // with the License. You may obtain a copy of the License at |
8 | | // |
9 | | // http://www.apache.org/licenses/LICENSE-2.0 |
10 | | // |
11 | | // Unless required by applicable law or agreed to in writing, |
12 | | // software distributed under the License is distributed on an |
13 | | // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
14 | | // KIND, either express or implied. See the License for the |
15 | | // specific language governing permissions and limitations |
16 | | // under the License. |
17 | | |
18 | | #include "format_v2/table/adbc_reader.h" |
19 | | |
20 | | #include <arrow-adbc/adbc.h> |
21 | | #include <arrow-adbc/adbc_driver_manager.h> |
22 | | #include <arrow/array/array_base.h> |
23 | | #include <arrow/c/bridge.h> |
24 | | #include <arrow/record_batch.h> |
25 | | |
26 | | #include <cstdint> |
27 | | #include <map> |
28 | | #include <memory> |
29 | | #include <string> |
30 | | #include <string_view> |
31 | | #include <utility> |
32 | | #include <vector> |
33 | | |
34 | | #include "common/cast_set.h" |
35 | | #include "common/check.h" |
36 | | #include "core/assert_cast.h" |
37 | | #include "core/block/block.h" |
38 | | #include "core/data_type/data_type.h" |
39 | | #include "core/data_type/data_type_array.h" |
40 | | #include "core/data_type/data_type_map.h" |
41 | | #include "core/data_type/data_type_nullable.h" |
42 | | #include "core/data_type/data_type_struct.h" |
43 | | #include "core/data_type/primitive_type.h" |
44 | | #include "core/data_type_serde/data_type_serde.h" |
45 | | #include "format/arrow/arrow_array_normalizer.h" |
46 | | #include "format/parquet/arrow_memory_pool.h" |
47 | | #include "format_v2/column_mapper.h" |
48 | | #include "format_v2/materialized_reader_util.h" |
49 | | #include "runtime/descriptors.h" |
50 | | #include "runtime/exec_env.h" |
51 | | #include "runtime/file_scan_profile.h" |
52 | | #include "runtime/runtime_state.h" |
53 | | #include "util/adbc_driver_registry.h" |
54 | | #include "util/string_util.h" |
55 | | #include "util/timezone_utils.h" |
56 | | #include "util/url_coding.h" |
57 | | |
58 | | namespace doris::format::adbc { |
59 | | namespace { |
60 | | |
61 | | // Keys of TTableFormatFileDesc.adbc_params. Kept next to the code that reads them so the FE-side |
62 | | // producer and this consumer stay diffable. |
63 | | constexpr const char* kParamDriverPath = "driver_path"; |
64 | | constexpr const char* kParamDriverEntrypoint = "driver_entrypoint"; |
65 | | constexpr const char* kParamUri = "uri"; |
66 | | constexpr const char* kParamUsername = "username"; |
67 | | constexpr const char* kParamPassword = "password"; |
68 | | constexpr const char* kParamQuerySql = "query_sql"; |
69 | | // Base64 of one opaque partition descriptor the driver produced on FE. Mutually exclusive with |
70 | | // kParamQuerySql: a range either runs a statement here or reads one partition of a statement the |
71 | | // source has already run. |
72 | | constexpr const char* kParamPartitionDescriptor = "partition_descriptor"; |
73 | | // Anything under this prefix is an ADBC option name in full (the prefix is part of the option name, |
74 | | // e.g. "adbc.connection.autocommit") and is handed to the driver untouched. |
75 | | constexpr std::string_view kAdbcOptionPrefix = "adbc."; |
76 | | |
77 | | const std::string* find_param(const std::map<std::string, std::string>& params, |
78 | 137 | const std::string& key) { |
79 | 137 | const auto it = params.find(key); |
80 | 137 | return it == params.end() ? nullptr : &it->second; |
81 | 137 | } |
82 | | |
83 | 29 | Status validate_adbc_range(const TFileRangeDesc& range) { |
84 | 29 | if (!range.__isset.table_format_params || |
85 | 29 | range.table_format_params.table_format_type != "adbc") { |
86 | 0 | return Status::InvalidArgument("ADBC reader requires the adbc table format"); |
87 | 0 | } |
88 | 29 | if (!range.table_format_params.__isset.adbc_params) { |
89 | 1 | return Status::InvalidArgument("ADBC reader requires adbc_params"); |
90 | 1 | } |
91 | 28 | const auto& params = range.table_format_params.adbc_params; |
92 | 55 | for (const auto* key : {kParamDriverPath, kParamUri}) { |
93 | 55 | const auto* value = find_param(params, key); |
94 | 55 | if (value == nullptr || value->empty()) { |
95 | 2 | return Status::InvalidArgument("ADBC reader requires a non-empty '{}' parameter", key); |
96 | 2 | } |
97 | 55 | } |
98 | 26 | const auto* query_sql = find_param(params, kParamQuerySql); |
99 | 26 | const auto* partition = find_param(params, kParamPartitionDescriptor); |
100 | 26 | const bool has_query = query_sql != nullptr && !query_sql->empty(); |
101 | 26 | const bool has_partition = partition != nullptr && !partition->empty(); |
102 | | // Not a defensive nicety: reading a partition means the source has ALREADY run the statement, so |
103 | | // a range carrying both would let this reader run it a second time depending on which branch it |
104 | | // happened to take. FE refuses to build such a range; this refuses to act on one. |
105 | 26 | if (has_query == has_partition) { |
106 | 3 | return Status::InvalidArgument( |
107 | 3 | "ADBC reader requires exactly one of '{}' and '{}', but the range carries {}", |
108 | 3 | kParamQuerySql, kParamPartitionDescriptor, has_query ? "both" : "neither"); |
109 | 3 | } |
110 | 23 | return Status::OK(); |
111 | 26 | } |
112 | | |
113 | | Status arrow_type_mismatch(const std::string& path, const arrow::DataType& arrow_type, |
114 | 3 | const DataTypePtr& doris_type) { |
115 | 3 | return Status::InvalidArgument( |
116 | 3 | "ADBC Arrow type mismatch at '{}': runtime type '{}' does not match cached Doris " |
117 | 3 | "type '{}'", |
118 | 3 | path, arrow_type.ToString(), doris_type->get_name()); |
119 | 3 | } |
120 | | |
121 | 0 | int timestamp_scale(arrow::TimeUnit::type unit) { |
122 | 0 | switch (unit) { |
123 | 0 | case arrow::TimeUnit::SECOND: |
124 | 0 | return 0; |
125 | 0 | case arrow::TimeUnit::MILLI: |
126 | 0 | return 3; |
127 | 0 | case arrow::TimeUnit::MICRO: |
128 | 0 | case arrow::TimeUnit::NANO: |
129 | 0 | return 6; |
130 | 0 | } |
131 | 0 | return -1; |
132 | 0 | } |
133 | | |
134 | | Status validate_arrow_type(const arrow::DataType& arrow_type, const DataTypePtr& doris_type, |
135 | 18 | const std::string& path) { |
136 | 18 | DORIS_CHECK(doris_type != nullptr); |
137 | 18 | const auto nested_doris_type = remove_nullable(doris_type); |
138 | 18 | const auto primitive_type = nested_doris_type->get_primitive_type(); |
139 | 18 | const auto mismatch = [&]() { return arrow_type_mismatch(path, arrow_type, doris_type); }; |
140 | | |
141 | 18 | switch (arrow_type.id()) { |
142 | 0 | case arrow::Type::BOOL: |
143 | 0 | return primitive_type == TYPE_BOOLEAN ? Status::OK() : mismatch(); |
144 | 0 | case arrow::Type::INT8: |
145 | 0 | return primitive_type == TYPE_TINYINT ? Status::OK() : mismatch(); |
146 | 0 | case arrow::Type::UINT8: |
147 | 0 | return primitive_type == TYPE_SMALLINT ? Status::OK() : mismatch(); |
148 | 0 | case arrow::Type::INT16: |
149 | 0 | return primitive_type == TYPE_SMALLINT ? Status::OK() : mismatch(); |
150 | 0 | case arrow::Type::UINT16: |
151 | 0 | return primitive_type == TYPE_INT ? Status::OK() : mismatch(); |
152 | 3 | case arrow::Type::INT32: |
153 | 3 | return primitive_type == TYPE_INT ? Status::OK() : mismatch(); |
154 | 0 | case arrow::Type::UINT32: |
155 | 0 | return primitive_type == TYPE_BIGINT ? Status::OK() : mismatch(); |
156 | 3 | case arrow::Type::INT64: |
157 | 3 | return primitive_type == TYPE_BIGINT ? Status::OK() : mismatch(); |
158 | 0 | case arrow::Type::UINT64: |
159 | 0 | return primitive_type == TYPE_LARGEINT ? Status::OK() : mismatch(); |
160 | 0 | case arrow::Type::HALF_FLOAT: |
161 | 2 | case arrow::Type::FLOAT: |
162 | 2 | return primitive_type == TYPE_FLOAT ? Status::OK() : mismatch(); |
163 | 1 | case arrow::Type::DOUBLE: |
164 | 1 | return primitive_type == TYPE_DOUBLE ? Status::OK() : mismatch(); |
165 | 2 | case arrow::Type::STRING: |
166 | 2 | case arrow::Type::BINARY: |
167 | 2 | case arrow::Type::FIXED_SIZE_BINARY: |
168 | 2 | return primitive_type == TYPE_STRING ? Status::OK() : mismatch(); |
169 | 0 | case arrow::Type::DATE32: |
170 | 0 | return primitive_type == TYPE_DATEV2 ? Status::OK() : mismatch(); |
171 | 0 | case arrow::Type::DATE64: |
172 | 0 | return primitive_type == TYPE_DATETIMEV2 && nested_doris_type->get_scale() == 3 |
173 | 0 | ? Status::OK() |
174 | 0 | : mismatch(); |
175 | 0 | case arrow::Type::TIMESTAMP: { |
176 | 0 | const auto& timestamp = static_cast<const arrow::TimestampType&>(arrow_type); |
177 | 0 | const bool zoned = !timestamp.timezone().empty(); |
178 | 0 | const auto expected = zoned ? TYPE_TIMESTAMPTZ : TYPE_DATETIMEV2; |
179 | 0 | return primitive_type == expected && |
180 | 0 | nested_doris_type->get_scale() == timestamp_scale(timestamp.unit()) |
181 | 0 | ? Status::OK() |
182 | 0 | : mismatch(); |
183 | 2 | } |
184 | 0 | case arrow::Type::DECIMAL128: |
185 | 0 | case arrow::Type::DECIMAL256: { |
186 | 0 | const auto& decimal = static_cast<const arrow::DecimalType&>(arrow_type); |
187 | 0 | return is_decimal(primitive_type) && |
188 | 0 | nested_doris_type->get_precision() == decimal.precision() && |
189 | 0 | nested_doris_type->get_scale() == decimal.scale() |
190 | 0 | ? Status::OK() |
191 | 0 | : mismatch(); |
192 | 0 | } |
193 | 4 | case arrow::Type::LIST: |
194 | 5 | case arrow::Type::LARGE_LIST: |
195 | 5 | case arrow::Type::FIXED_SIZE_LIST: { |
196 | 5 | if (primitive_type != TYPE_ARRAY) { |
197 | 0 | return mismatch(); |
198 | 0 | } |
199 | 5 | const auto& list = static_cast<const arrow::BaseListType&>(arrow_type); |
200 | 5 | const auto& array = assert_cast<const DataTypeArray&>(*nested_doris_type); |
201 | 5 | return validate_arrow_type(*list.value_type(), array.get_nested_type(), path + ".element"); |
202 | 5 | } |
203 | 2 | case arrow::Type::STRUCT: { |
204 | 2 | if (primitive_type != TYPE_STRUCT) { |
205 | 0 | return mismatch(); |
206 | 0 | } |
207 | 2 | const auto& arrow_struct = static_cast<const arrow::StructType&>(arrow_type); |
208 | 2 | const auto& doris_struct = assert_cast<const DataTypeStruct&>(*nested_doris_type); |
209 | 2 | if (cast_set<size_t>(arrow_struct.num_fields()) != doris_struct.get_elements().size()) { |
210 | 0 | return mismatch(); |
211 | 0 | } |
212 | 4 | for (int field_idx = 0; field_idx < arrow_struct.num_fields(); ++field_idx) { |
213 | 3 | const auto& arrow_field = arrow_struct.field(field_idx); |
214 | 3 | const auto& doris_name = doris_struct.get_element_name(field_idx); |
215 | 3 | if (to_lower(arrow_field->name()) != to_lower(doris_name)) { |
216 | 1 | return Status::InvalidArgument( |
217 | 1 | "ADBC Arrow field mismatch at '{}': runtime field '{}' does not match " |
218 | 1 | "cached field '{}' at ordinal {}", |
219 | 1 | path, arrow_field->name(), doris_name, field_idx); |
220 | 1 | } |
221 | 2 | RETURN_IF_ERROR(validate_arrow_type(*arrow_field->type(), |
222 | 2 | doris_struct.get_element(field_idx), |
223 | 2 | path + "." + doris_name)); |
224 | 2 | } |
225 | 1 | return Status::OK(); |
226 | 2 | } |
227 | 0 | case arrow::Type::MAP: { |
228 | 0 | if (primitive_type != TYPE_MAP) { |
229 | 0 | return mismatch(); |
230 | 0 | } |
231 | 0 | const auto& arrow_map = static_cast<const arrow::MapType&>(arrow_type); |
232 | 0 | const auto& doris_map = assert_cast<const DataTypeMap&>(*nested_doris_type); |
233 | 0 | RETURN_IF_ERROR(validate_arrow_type(*arrow_map.key_type(), doris_map.get_key_type(), |
234 | 0 | path + ".key")); |
235 | 0 | return validate_arrow_type(*arrow_map.item_type(), doris_map.get_value_type(), |
236 | 0 | path + ".value"); |
237 | 0 | } |
238 | 0 | default: |
239 | 0 | return mismatch(); |
240 | 18 | } |
241 | 18 | } |
242 | | |
243 | | // Drivers allocate the strings inside AdbcError, so every populated error has to be released. |
244 | | class AdbcErrorGuard { |
245 | | public: |
246 | 10 | AdbcErrorGuard() = default; |
247 | 10 | ~AdbcErrorGuard() { reset(); } |
248 | | AdbcErrorGuard(const AdbcErrorGuard&) = delete; |
249 | | AdbcErrorGuard& operator=(const AdbcErrorGuard&) = delete; |
250 | | |
251 | 37 | AdbcError* get() { return &_error; } |
252 | | |
253 | 2 | std::string take_message() { |
254 | 2 | std::string message = _error.message != nullptr ? _error.message : ""; |
255 | 2 | reset(); |
256 | 2 | return message; |
257 | 2 | } |
258 | | |
259 | 47 | void reset() { |
260 | 47 | if (_error.release != nullptr) { |
261 | 1 | _error.release(&_error); |
262 | 1 | } |
263 | 47 | _error = ADBC_ERROR_INIT; |
264 | 47 | } |
265 | | |
266 | | private: |
267 | | AdbcError _error = ADBC_ERROR_INIT; |
268 | | }; |
269 | | |
270 | 2 | Status adbc_call_status(const char* what, AdbcStatusCode code, AdbcErrorGuard& error) { |
271 | 2 | const std::string message = error.take_message(); |
272 | | // AdbcStatusCode is a uint8_t, so spell out the name as well as the number. |
273 | 2 | return Status::InternalError("ADBC: {} failed ({}, code {}): {}", what, |
274 | 2 | AdbcStatusCodeMessage(code), static_cast<int>(code), |
275 | 2 | message.empty() ? "driver reported no message" : message); |
276 | 2 | } |
277 | | |
278 | | #define RETURN_IF_ADBC_ERROR(expr, what, error) \ |
279 | 27 | do { \ |
280 | 27 | const AdbcStatusCode adbc_call_code = (expr); \ |
281 | 27 | if (adbc_call_code != ADBC_STATUS_OK) { \ |
282 | 2 | return adbc_call_status((what), adbc_call_code, (error)); \ |
283 | 2 | } \ |
284 | 27 | (error).reset(); \ |
285 | 25 | } while (0) |
286 | | |
287 | | // The production stream: one ADBC database/connection/statement per scan range. |
288 | | // |
289 | | // P0 keeps them un-pooled on purpose. Reusing databases across ranges is a throughput optimization |
290 | | // that only pays off once multiple partitions run concurrently, and an unverifiable caching layer |
291 | | // added now would only obscure the functional path. |
292 | | class RealAdbcStream final : public AdbcStream { |
293 | | public: |
294 | 5 | explicit RealAdbcStream(const TFileRangeDesc& range) : _range(range) {} |
295 | 5 | ~RealAdbcStream() override { static_cast<void>(close()); } |
296 | | |
297 | 5 | Status open() { |
298 | 5 | RETURN_IF_ERROR(validate_adbc_range(_range)); |
299 | 5 | const auto& params = _range.table_format_params.adbc_params; |
300 | 5 | const std::string& driver_path = *find_param(params, kParamDriverPath); |
301 | 5 | const std::string& uri = *find_param(params, kParamUri); |
302 | 5 | const auto* entrypoint = find_param(params, kParamDriverEntrypoint); |
303 | | // validate_adbc_range has established that exactly one of these is present. |
304 | 5 | const auto* partition = find_param(params, kParamPartitionDescriptor); |
305 | | |
306 | 5 | RETURN_IF_ERROR(AdbcDriverRegistry::instance().get_or_load( |
307 | 5 | driver_path, entrypoint != nullptr ? *entrypoint : std::string(), &_driver)); |
308 | 4 | DORIS_CHECK(_driver != nullptr); |
309 | | |
310 | 4 | AdbcErrorGuard error; |
311 | 4 | RETURN_IF_ADBC_ERROR(_driver->DatabaseNew(&_database, error.get()), "DatabaseNew", error); |
312 | 4 | _database_created = true; |
313 | 4 | RETURN_IF_ERROR(_set_database_options(params, uri, error)); |
314 | 4 | RETURN_IF_ADBC_ERROR(_driver->DatabaseInit(&_database, error.get()), "DatabaseInit", error); |
315 | | |
316 | 4 | RETURN_IF_ADBC_ERROR(_driver->ConnectionNew(&_connection, error.get()), "ConnectionNew", |
317 | 4 | error); |
318 | 4 | _connection_created = true; |
319 | 4 | RETURN_IF_ADBC_ERROR(_driver->ConnectionInit(&_connection, &_database, error.get()), |
320 | 4 | "ConnectionInit", error); |
321 | | |
322 | 4 | if (partition != nullptr && !partition->empty()) { |
323 | 2 | RETURN_IF_ERROR(_read_partition(*partition, error)); |
324 | 2 | } else { |
325 | 2 | RETURN_IF_ERROR(_execute_query(*find_param(params, kParamQuerySql), error)); |
326 | 2 | } |
327 | | |
328 | | // Before Arrow ever sees it: the driver's stream may not clear its release callback, and |
329 | | // Arrow aborts the process when that happens. Both branches need it -- the Flight SQL |
330 | | // driver's ReadPartition stream breaks the contract exactly like its ExecuteQuery one. |
331 | 1 | enforce_stream_release_contract(&_c_stream); |
332 | | |
333 | 1 | auto reader = arrow::ImportRecordBatchReader(&_c_stream); |
334 | 1 | if (!reader.ok()) { |
335 | 0 | return Status::InternalError("ADBC: failed to import the result stream: {}", |
336 | 0 | reader.status().ToString()); |
337 | 0 | } |
338 | | // ImportRecordBatchReader moves the stream's contents; the reader owns it from here. |
339 | 1 | _reader = reader.MoveValueUnsafe(); |
340 | 1 | return Status::OK(); |
341 | 1 | } |
342 | | |
343 | 1 | Status next(std::shared_ptr<arrow::RecordBatch>* batch) override { |
344 | 1 | DORIS_CHECK(batch != nullptr); |
345 | 1 | if (_reader == nullptr) { |
346 | 0 | return Status::InternalError("ADBC: result stream is not open"); |
347 | 0 | } |
348 | 1 | std::shared_ptr<arrow::RecordBatch> next_batch; |
349 | 1 | const auto status = _reader->ReadNext(&next_batch); |
350 | 1 | if (!status.ok()) { |
351 | 0 | return Status::InternalError("ADBC: failed to read the next batch: {}", |
352 | 0 | status.ToString()); |
353 | 0 | } |
354 | 1 | *batch = std::move(next_batch); |
355 | 1 | return Status::OK(); |
356 | 1 | } |
357 | | |
358 | 6 | Status close() override { |
359 | 6 | Status result = Status::OK(); |
360 | | // Release in reverse order of creation. The reader owns the imported stream, so it has to |
361 | | // go before the statement that produced it. |
362 | 6 | _reader.reset(); |
363 | 6 | if (_c_stream.release != nullptr) { |
364 | | // Only reachable when the import itself failed; nothing else owns the stream then. |
365 | 0 | _c_stream.release(&_c_stream); |
366 | 0 | _c_stream = {}; |
367 | 0 | } |
368 | 6 | AdbcErrorGuard error; |
369 | 6 | if (_statement_created) { |
370 | 2 | const auto code = _driver->StatementRelease(&_statement, error.get()); |
371 | 2 | if (code != ADBC_STATUS_OK && result.ok()) { |
372 | 0 | result = adbc_call_status("StatementRelease", code, error); |
373 | 0 | } |
374 | 2 | error.reset(); |
375 | 2 | _statement_created = false; |
376 | 2 | } |
377 | 6 | if (_connection_created) { |
378 | 4 | const auto code = _driver->ConnectionRelease(&_connection, error.get()); |
379 | 4 | if (code != ADBC_STATUS_OK && result.ok()) { |
380 | 0 | result = adbc_call_status("ConnectionRelease", code, error); |
381 | 0 | } |
382 | 4 | error.reset(); |
383 | 4 | _connection_created = false; |
384 | 4 | } |
385 | 6 | if (_database_created) { |
386 | 4 | const auto code = _driver->DatabaseRelease(&_database, error.get()); |
387 | 4 | if (code != ADBC_STATUS_OK && result.ok()) { |
388 | 0 | result = adbc_call_status("DatabaseRelease", code, error); |
389 | 0 | } |
390 | 4 | error.reset(); |
391 | 4 | _database_created = false; |
392 | 4 | } |
393 | | // _driver itself is owned by AdbcDriverRegistry and is never released. |
394 | 6 | return result; |
395 | 6 | } |
396 | | |
397 | | private: |
398 | | // Runs the statement FE generated. One statement per range, so this range is the whole query. |
399 | 2 | Status _execute_query(const std::string& query_sql, AdbcErrorGuard& error) { |
400 | 2 | RETURN_IF_ADBC_ERROR(_driver->StatementNew(&_connection, &_statement, error.get()), |
401 | 2 | "StatementNew", error); |
402 | 2 | _statement_created = true; |
403 | 2 | RETURN_IF_ADBC_ERROR( |
404 | 2 | _driver->StatementSetSqlQuery(&_statement, query_sql.c_str(), error.get()), |
405 | 2 | "StatementSetSqlQuery", error); |
406 | 2 | int64_t rows_affected = -1; |
407 | 2 | RETURN_IF_ADBC_ERROR(_driver->StatementExecuteQuery(&_statement, &_c_stream, &rows_affected, |
408 | 2 | error.get()), |
409 | 2 | "StatementExecuteQuery", error); |
410 | 1 | return Status::OK(); |
411 | 2 | } |
412 | | |
413 | | // Reads one partition of a query FE already had the source execute. No statement is created: |
414 | | // ADBC reads a partition off a connection, and the whole point is that this can happen on a |
415 | | // different machine from the one that planned it. |
416 | 2 | Status _read_partition(const std::string& base64_descriptor, AdbcErrorGuard& error) { |
417 | 2 | std::string descriptor; |
418 | 2 | if (!base64_decode(base64_descriptor, &descriptor)) { |
419 | 1 | return Status::InvalidArgument("ADBC: the '{}' parameter is not valid base64", |
420 | 1 | kParamPartitionDescriptor); |
421 | 1 | } |
422 | 1 | RETURN_IF_ADBC_ERROR( |
423 | 1 | _driver->ConnectionReadPartition( |
424 | 1 | &_connection, reinterpret_cast<const uint8_t*>(descriptor.data()), |
425 | 1 | descriptor.size(), &_c_stream, error.get()), |
426 | 1 | "ConnectionReadPartition", error); |
427 | 0 | return Status::OK(); |
428 | 1 | } |
429 | | |
430 | | Status _set_database_options(const std::map<std::string, std::string>& params, |
431 | 4 | const std::string& uri, AdbcErrorGuard& error) { |
432 | 4 | RETURN_IF_ADBC_ERROR( |
433 | 4 | _driver->DatabaseSetOption(&_database, ADBC_OPTION_URI, uri.c_str(), error.get()), |
434 | 4 | "DatabaseSetOption(uri)", error); |
435 | 8 | for (const auto* key : {kParamUsername, kParamPassword}) { |
436 | 8 | const auto* value = find_param(params, key); |
437 | 8 | if (value == nullptr || value->empty()) { |
438 | 8 | continue; |
439 | 8 | } |
440 | 0 | RETURN_IF_ADBC_ERROR( |
441 | 0 | _driver->DatabaseSetOption(&_database, key, value->c_str(), error.get()), |
442 | 0 | "DatabaseSetOption(credentials)", error); |
443 | 0 | } |
444 | 12 | for (const auto& [key, value] : params) { |
445 | 12 | if (!key.starts_with(kAdbcOptionPrefix)) { |
446 | 12 | continue; |
447 | 12 | } |
448 | 0 | RETURN_IF_ADBC_ERROR( |
449 | 0 | _driver->DatabaseSetOption(&_database, key.c_str(), value.c_str(), error.get()), |
450 | 0 | "DatabaseSetOption(passthrough)", error); |
451 | 0 | } |
452 | 4 | return Status::OK(); |
453 | 4 | } |
454 | | |
455 | | const TFileRangeDesc _range; |
456 | | const AdbcDriver* _driver = nullptr; |
457 | | AdbcDatabase _database {}; |
458 | | AdbcConnection _connection {}; |
459 | | AdbcStatement _statement {}; |
460 | | ArrowArrayStream _c_stream {}; |
461 | | std::shared_ptr<arrow::RecordBatchReader> _reader; |
462 | | bool _database_created = false; |
463 | | bool _connection_created = false; |
464 | | bool _statement_created = false; |
465 | | }; |
466 | | |
467 | 5 | Status create_real_adbc_stream(const TFileRangeDesc& range, std::unique_ptr<AdbcStream>* out) { |
468 | 5 | DORIS_CHECK(out != nullptr); |
469 | 5 | auto stream = std::make_unique<RealAdbcStream>(range); |
470 | 5 | RETURN_IF_ERROR(stream->open()); |
471 | 1 | *out = std::move(stream); |
472 | 1 | return Status::OK(); |
473 | 5 | } |
474 | | |
475 | | ColumnDefinition adbc_child_definition(const std::string& name, DataTypePtr type, int32_t local_id); |
476 | | |
477 | | // Mirrors synthesize_remote_doris_children in remote_doris_reader.cpp. Both readers expose table |
478 | | // slots as file columns, so complex columns still need structural children for TableColumnMapper. |
479 | | // Kept separate rather than shared to avoid reshaping the already-shipped remote_doris reader. |
480 | 26 | std::vector<ColumnDefinition> synthesize_adbc_children(const DataTypePtr& type) { |
481 | 26 | std::vector<ColumnDefinition> children; |
482 | 26 | DORIS_CHECK(type != nullptr); |
483 | 26 | const auto nested_type = remove_nullable(type); |
484 | 26 | switch (nested_type->get_primitive_type()) { |
485 | 7 | case TYPE_ARRAY: { |
486 | 7 | const auto* array_type = assert_cast<const DataTypeArray*>(nested_type.get()); |
487 | 7 | children.push_back(adbc_child_definition("element", array_type->get_nested_type(), 0)); |
488 | 7 | break; |
489 | 0 | } |
490 | 0 | case TYPE_MAP: { |
491 | 0 | const auto* map_type = assert_cast<const DataTypeMap*>(nested_type.get()); |
492 | 0 | children.push_back(adbc_child_definition("key", map_type->get_key_type(), 0)); |
493 | 0 | children.push_back(adbc_child_definition("value", map_type->get_value_type(), 1)); |
494 | 0 | break; |
495 | 0 | } |
496 | 2 | case TYPE_STRUCT: { |
497 | 2 | const auto* struct_type = assert_cast<const DataTypeStruct*>(nested_type.get()); |
498 | 2 | children.reserve(struct_type->get_elements().size()); |
499 | 6 | for (size_t idx = 0; idx < struct_type->get_elements().size(); ++idx) { |
500 | 4 | children.push_back(adbc_child_definition(struct_type->get_element_name(idx), |
501 | 4 | struct_type->get_element(idx), |
502 | 4 | cast_set<int32_t>(idx))); |
503 | 4 | } |
504 | 2 | break; |
505 | 0 | } |
506 | 17 | default: |
507 | 17 | break; |
508 | 26 | } |
509 | 26 | return children; |
510 | 26 | } |
511 | | |
512 | | ColumnDefinition adbc_child_definition(const std::string& name, DataTypePtr type, |
513 | 11 | int32_t local_id) { |
514 | 11 | ColumnDefinition child; |
515 | 11 | child.identifier = Field::create_field<TYPE_STRING>(name); |
516 | 11 | child.local_id = local_id; |
517 | 11 | child.name = name; |
518 | 11 | child.type = std::move(type); |
519 | 11 | child.children = synthesize_adbc_children(child.type); |
520 | 11 | return child; |
521 | 11 | } |
522 | | |
523 | | // A stream that forwards everything to the driver's and, on release, does the one thing some |
524 | | // drivers forget: clear its own release callback. Heap-allocated because Arrow keeps only the |
525 | | // ArrowArrayStream it was handed, and the delegate has to outlive this function. |
526 | | struct DelegatingStream { |
527 | | ArrowArrayStream inner; |
528 | | }; |
529 | | |
530 | 1 | int delegating_get_schema(ArrowArrayStream* self, ArrowSchema* out) { |
531 | 1 | auto& inner = static_cast<DelegatingStream*>(self->private_data)->inner; |
532 | 1 | return inner.get_schema(&inner, out); |
533 | 1 | } |
534 | | |
535 | 2 | int delegating_get_next(ArrowArrayStream* self, ArrowArray* out) { |
536 | 2 | auto& inner = static_cast<DelegatingStream*>(self->private_data)->inner; |
537 | 2 | return inner.get_next(&inner, out); |
538 | 2 | } |
539 | | |
540 | 1 | const char* delegating_get_last_error(ArrowArrayStream* self) { |
541 | 1 | auto& inner = static_cast<DelegatingStream*>(self->private_data)->inner; |
542 | 1 | return inner.get_last_error != nullptr ? inner.get_last_error(&inner) : nullptr; |
543 | 1 | } |
544 | | |
545 | 3 | void delegating_release(ArrowArrayStream* self) { |
546 | 3 | auto* delegate = static_cast<DelegatingStream*>(self->private_data); |
547 | 3 | if (delegate->inner.release != nullptr) { |
548 | 3 | delegate->inner.release(&delegate->inner); |
549 | 3 | } |
550 | 3 | delete delegate; |
551 | 3 | self->private_data = nullptr; |
552 | | // What the driver failed to do, and what Arrow aborts the process over. |
553 | 3 | self->release = nullptr; |
554 | 3 | } |
555 | | |
556 | | } // namespace |
557 | | |
558 | 4 | void enforce_stream_release_contract(ArrowArrayStream* stream) { |
559 | 4 | DORIS_CHECK(stream != nullptr); |
560 | 4 | if (stream->release == nullptr) { |
561 | | // Already released; nothing to delegate to, and wrapping it would hand Arrow a stream |
562 | | // whose callbacks dereference a released delegate. |
563 | 1 | return; |
564 | 1 | } |
565 | 3 | auto* delegate = new DelegatingStream {.inner = *stream}; |
566 | 3 | *stream = ArrowArrayStream {.get_schema = delegating_get_schema, |
567 | 3 | .get_next = delegating_get_next, |
568 | 3 | .get_last_error = delegating_get_last_error, |
569 | 3 | .release = delegating_release, |
570 | 3 | .private_data = delegate}; |
571 | 3 | } |
572 | | |
573 | | AdbcFileReader::AdbcFileReader(std::shared_ptr<io::FileSystemProperties>& system_properties, |
574 | | std::unique_ptr<io::FileDescription>& file_description, |
575 | | std::shared_ptr<io::IOContext> io_ctx, RuntimeProfile* profile, |
576 | | const TFileRangeDesc& range, |
577 | | const std::vector<SlotDescriptor*>& file_slot_descs, |
578 | | AdbcStreamFactory stream_factory) |
579 | 24 | : FileReader(system_properties, file_description, std::move(io_ctx), profile), |
580 | 24 | _range(range), |
581 | 24 | _file_slot_descs(file_slot_descs), |
582 | 24 | _stream_factory(std::move(stream_factory)) { |
583 | 24 | TimezoneUtils::find_cctz_time_zone(TimezoneUtils::default_time_zone, _ctz); |
584 | 24 | } |
585 | | |
586 | 24 | AdbcFileReader::~AdbcFileReader() { |
587 | 24 | static_cast<void>(close()); |
588 | 24 | } |
589 | | |
590 | 23 | void AdbcFileReader::_init_profile() { |
591 | 23 | if (_profile == nullptr) { |
592 | 0 | return; |
593 | 0 | } |
594 | 23 | const auto hierarchy = file_scan_profile::ensure_hierarchy(_profile); |
595 | 23 | _io_time = hierarchy.io; |
596 | 23 | static const char* adbc_profile = "AdbcFileReader"; |
597 | 23 | _total_time = |
598 | 23 | ADD_CHILD_TIMER_WITH_LEVEL(_profile, adbc_profile, file_scan_profile::FILE_READER, 1); |
599 | 23 | _open_stream_time = ADD_CHILD_TIMER_WITH_LEVEL(_profile, "AdbcOpenStreamTime", adbc_profile, 1); |
600 | 23 | _next_batch_time = ADD_CHILD_TIMER_WITH_LEVEL(_profile, "AdbcNextBatchTime", adbc_profile, 1); |
601 | 23 | _normalize_time = ADD_CHILD_TIMER_WITH_LEVEL(_profile, "AdbcNormalizeTime", adbc_profile, 1); |
602 | 23 | _materialize_time = |
603 | 23 | ADD_CHILD_TIMER_WITH_LEVEL(_profile, "AdbcMaterializeTime", adbc_profile, 1); |
604 | 23 | _filter_time = ADD_CHILD_TIMER_WITH_LEVEL(_profile, "AdbcFilterTime", adbc_profile, 1); |
605 | 23 | } |
606 | | |
607 | 23 | Status AdbcFileReader::init(RuntimeState* state) { |
608 | 23 | _init_profile(); |
609 | 23 | SCOPED_TIMER(_total_time); |
610 | 23 | _runtime_state = state; |
611 | 23 | RETURN_IF_ERROR(validate_adbc_range(_range)); |
612 | 17 | RETURN_IF_ERROR(_build_col_name_to_file_id()); |
613 | 17 | _eof = false; |
614 | 17 | return Status::OK(); |
615 | 17 | } |
616 | | |
617 | 13 | Status AdbcFileReader::get_schema(std::vector<ColumnDefinition>* file_schema) const { |
618 | 13 | SCOPED_TIMER(_total_time); |
619 | 13 | DORIS_CHECK(file_schema != nullptr); |
620 | 13 | file_schema->clear(); |
621 | 13 | file_schema->reserve(_file_slot_descs.size()); |
622 | 28 | for (size_t idx = 0; idx < _file_slot_descs.size(); ++idx) { |
623 | 15 | const auto* slot = _file_slot_descs[idx]; |
624 | 15 | DORIS_CHECK(slot != nullptr); |
625 | 15 | file_schema->push_back({ |
626 | 15 | .identifier = Field::create_field<TYPE_INT>(cast_set<int32_t>(idx)), |
627 | 15 | .local_id = cast_set<int32_t>(idx), |
628 | 15 | .name = slot->col_name(), |
629 | 15 | .type = slot->type(), |
630 | 15 | .children = synthesize_adbc_children(slot->type()), |
631 | 15 | }); |
632 | 15 | } |
633 | 13 | return Status::OK(); |
634 | 13 | } |
635 | | |
636 | | std::unique_ptr<TableColumnMapper> AdbcFileReader::create_column_mapper( |
637 | 2 | TableColumnMapperOptions options) const { |
638 | | // ADBC streams return complete Arrow roots, so TableReader must own every nested projection and |
639 | | // reorder after the full source shape has been materialized. |
640 | 2 | return std::make_unique<MaterializedColumnMapper>(std::move(options)); |
641 | 2 | } |
642 | | |
643 | 17 | Status AdbcFileReader::open(std::shared_ptr<FileScanRequest> request) { |
644 | 17 | SCOPED_TIMER(_total_time); |
645 | 17 | SCOPED_TIMER(_open_stream_time); |
646 | 17 | RETURN_IF_ERROR(FileReader::open(std::move(request))); |
647 | 17 | RETURN_IF_ERROR(_open_stream()); |
648 | 13 | _eof = false; |
649 | 13 | return Status::OK(); |
650 | 17 | } |
651 | | |
652 | 15 | Status AdbcFileReader::get_block(Block* file_block, size_t* rows, bool* eof) { |
653 | 15 | SCOPED_TIMER(_total_time); |
654 | 15 | DORIS_CHECK(file_block != nullptr); |
655 | 15 | DORIS_CHECK(rows != nullptr); |
656 | 15 | DORIS_CHECK(eof != nullptr); |
657 | 15 | if (_stream == nullptr) { |
658 | 0 | return Status::InternalError("ADBC reader is not open"); |
659 | 0 | } |
660 | 15 | if (_io_ctx != nullptr && _io_ctx->should_stop) { |
661 | | // Observe cancellation before entering a potentially blocking driver read. |
662 | 0 | RETURN_IF_ERROR(close()); |
663 | 0 | *rows = 0; |
664 | 0 | *eof = true; |
665 | 0 | return Status::OK(); |
666 | 0 | } |
667 | | |
668 | 15 | *rows = 0; |
669 | 15 | *eof = false; |
670 | 15 | std::shared_ptr<arrow::RecordBatch> batch; |
671 | 15 | { |
672 | 15 | SCOPED_TIMER(_io_time); |
673 | 15 | SCOPED_TIMER(_next_batch_time); |
674 | 15 | RETURN_IF_ERROR(_stream->next(&batch)); |
675 | 15 | } |
676 | 15 | if (batch == nullptr) { |
677 | 2 | *eof = true; |
678 | 2 | _eof = true; |
679 | 2 | return Status::OK(); |
680 | 2 | } |
681 | | |
682 | 13 | { |
683 | 13 | SCOPED_TIMER(_materialize_time); |
684 | 13 | RETURN_IF_ERROR(_materialize_record_batch(*batch, file_block, rows)); |
685 | 13 | } |
686 | 6 | _record_scan_rows(cast_set<int64_t>(*rows)); |
687 | 6 | { |
688 | 6 | SCOPED_TIMER(_filter_time); |
689 | 6 | RETURN_IF_ERROR( |
690 | 6 | apply_materialized_reader_filters(_request.get(), _io_ctx.get(), file_block, rows)); |
691 | 6 | } |
692 | 6 | return Status::OK(); |
693 | 6 | } |
694 | | |
695 | 34 | Status AdbcFileReader::close() { |
696 | 34 | SCOPED_TIMER(_total_time); |
697 | 34 | if (_stream != nullptr) { |
698 | 13 | RETURN_IF_ERROR(_stream->close()); |
699 | 13 | _stream.reset(); |
700 | 13 | } |
701 | 34 | _request.reset(); |
702 | 34 | _eof = true; |
703 | 34 | return Status::OK(); |
704 | 34 | } |
705 | | |
706 | 17 | Status AdbcFileReader::_open_stream() { |
707 | 17 | DORIS_CHECK(_stream == nullptr); |
708 | 17 | if (_stream_factory) { |
709 | 12 | RETURN_IF_ERROR(_stream_factory(_range, &_stream)); |
710 | 12 | } else { |
711 | 5 | RETURN_IF_ERROR(create_real_adbc_stream(_range, &_stream)); |
712 | 5 | } |
713 | 13 | DORIS_CHECK(_stream != nullptr); |
714 | 13 | return Status::OK(); |
715 | 17 | } |
716 | | |
717 | | Status AdbcFileReader::_materialize_record_batch(const arrow::RecordBatch& batch, Block* file_block, |
718 | 13 | size_t* rows) const { |
719 | 13 | DORIS_CHECK(file_block != nullptr); |
720 | 13 | DORIS_CHECK(rows != nullptr); |
721 | 13 | if (_request == nullptr) { |
722 | 0 | return Status::InternalError("ADBC reader is not open"); |
723 | 0 | } |
724 | | |
725 | 13 | if (_col_name_to_file_id.empty()) { |
726 | | // A pushed-down COUNT(*) projects no columns at all: the scan wants rows counted, no values. |
727 | | // Counting here rather than falling into the loop below is not an optimization -- every column |
728 | | // the source returns is unrequested by definition in this state, so the loop's unknown-column |
729 | | // check would reject the first one and fail a query that asked for nothing but a number. |
730 | | // FE sends a one-constant-column statement for this case, so the batch is narrow. |
731 | | // |
732 | | // Only the empty case is special-cased. An unrequested column arriving alongside requested ones |
733 | | // still fails: that means FE and this reader disagree about the projection, and it is the one |
734 | | // signal that the disagreement exists. |
735 | 1 | *rows = cast_set<size_t>(batch.num_rows()); |
736 | 1 | return Status::OK(); |
737 | 1 | } |
738 | | |
739 | 12 | ArrowMemoryPool<> local_arrow_pool; |
740 | 12 | arrow::MemoryPool* arrow_pool = ExecEnv::GetInstance()->arrow_memory_pool(); |
741 | 12 | if (arrow_pool == nullptr) { |
742 | | // Embedded and unit-test runtimes may omit ExecEnv memory initialization; keep conversions |
743 | | // on Doris' tracked allocator instead of falling back to Arrow's untracked default pool. |
744 | 12 | arrow_pool = &local_arrow_pool; |
745 | 12 | } |
746 | 12 | std::vector<bool> materialized_columns(file_block->columns(), false); |
747 | 20 | for (int arrow_idx = 0; arrow_idx < batch.num_columns(); ++arrow_idx) { |
748 | 15 | const std::string& column_name = batch.schema()->field(arrow_idx)->name(); |
749 | 15 | const auto file_id_it = _col_name_to_file_id.find(column_name); |
750 | 15 | if (file_id_it == _col_name_to_file_id.end()) { |
751 | 1 | return Status::InternalError("ADBC source returned unknown column {}", column_name); |
752 | 1 | } |
753 | 14 | const auto block_position_it = _request->local_positions.find(file_id_it->second); |
754 | 14 | if (block_position_it == _request->local_positions.end()) { |
755 | 0 | continue; |
756 | 0 | } |
757 | 14 | std::shared_ptr<arrow::Array> array; |
758 | 14 | { |
759 | 14 | SCOPED_TIMER(_normalize_time); |
760 | 14 | RETURN_IF_ERROR(normalize_arrow_array(batch.column(arrow_idx), arrow_pool, &array)); |
761 | 14 | } |
762 | 14 | RETURN_IF_ERROR(_materialize_arrow_column(column_name, array, batch.num_rows(), |
763 | 14 | file_id_it->second, block_position_it->second, |
764 | 14 | file_block)); |
765 | 8 | materialized_columns[block_position_it->second.value()] = true; |
766 | 8 | } |
767 | | |
768 | 8 | for (const auto& [file_column_id, block_position] : _request->local_positions) { |
769 | 8 | if (block_position.value() >= materialized_columns.size()) { |
770 | 0 | return Status::InternalError( |
771 | 0 | "ADBC requested block position {} out of range, block columns {}", |
772 | 0 | block_position.value(), materialized_columns.size()); |
773 | 0 | } |
774 | 8 | if (!materialized_columns[block_position.value()]) { |
775 | 0 | return Status::InternalError("ADBC source did not return requested file column id {}", |
776 | 0 | file_column_id.value()); |
777 | 0 | } |
778 | 8 | } |
779 | | |
780 | 5 | *rows = cast_set<size_t>(batch.num_rows()); |
781 | 5 | return Status::OK(); |
782 | 5 | } |
783 | | |
784 | | Status AdbcFileReader::_materialize_arrow_column(const std::string& column_name, |
785 | | const std::shared_ptr<arrow::Array>& array, |
786 | | int64_t num_rows, LocalColumnId file_column_id, |
787 | | const LocalIndex& block_position, |
788 | 14 | Block* file_block) const { |
789 | 14 | DORIS_CHECK(file_block != nullptr); |
790 | 14 | DORIS_CHECK(array != nullptr); |
791 | 14 | if (block_position.value() >= file_block->columns()) { |
792 | 0 | return Status::InternalError("ADBC block position {} out of range, block columns {}", |
793 | 0 | block_position.value(), file_block->columns()); |
794 | 0 | } |
795 | 14 | const auto& target_type = file_block->get_by_position(block_position.value()).type; |
796 | | |
797 | | // The cached FE target remains authoritative when a source schema changes mid-cache-window; |
798 | | // otherwise SerDes without a null map can silently turn source nulls into default values. |
799 | 14 | if (array->null_count() > 0 && !target_type->is_nullable()) { |
800 | 2 | return Status::InternalError( |
801 | 2 | "ADBC Arrow column '{}' contains {} null rows for non-nullable Doris type {}", |
802 | 2 | column_name, array->null_count(), target_type->get_name()); |
803 | 2 | } |
804 | | |
805 | | // An all-null column arrives with a type that says nothing about the column. |
806 | | // |
807 | | // A source that infers Arrow types from the VALUES it returns -- rather than from the declared |
808 | | // column type -- has nothing to infer from when every value in the result is null, and picks |
809 | | // whatever its default is. Measured against the SQLite driver: the same TEXT column comes back |
810 | | // as utf8 for `SELECT id, name FROM t1` and as int64 for |
811 | | // `SELECT id, name FROM t1 WHERE name IS NULL`, purely because the filter left only nulls. |
812 | | // Handing that to the serde fails with "Unsupported arrow type for string column: 9", and no |
813 | | // amount of care on the FE side avoids it: FE cannot know in advance which rows will survive. |
814 | | // |
815 | | // N nulls are what this array means whatever type it claims, so materialize them directly. The |
816 | | // check is narrow on purpose -- a column with even one non-null value keeps its real type and |
817 | | // still fails loudly on a genuine mismatch, which is the signal that FE and the source disagree |
818 | | // about the schema rather than about one result set. |
819 | | // |
820 | | // Only for a nullable target: substituting defaults into a NOT NULL column would turn a source |
821 | | // that wrongly sent nulls into silently wrong data, so that keeps failing in the serde. |
822 | 12 | if (array->null_count() == array->length() && target_type->is_nullable()) { |
823 | 1 | auto columns_guard = file_block->mutate_columns_scoped(); |
824 | 1 | auto& columns = columns_guard.mutable_columns(); |
825 | 1 | columns[block_position.value()]->insert_many_defaults(cast_set<size_t>(num_rows)); |
826 | 1 | return Status::OK(); |
827 | 1 | } |
828 | | |
829 | | // Validate the whole logical shape before acquiring mutable block columns. Struct SerDes map |
830 | | // children by ordinal, so a stale cache entry can otherwise produce plausible but wrong values. |
831 | 11 | RETURN_IF_ERROR(validate_arrow_type(*array->type(), target_type, column_name)); |
832 | | |
833 | 7 | auto columns_guard = file_block->mutate_columns_scoped(); |
834 | 7 | auto& columns = columns_guard.mutable_columns(); |
835 | 7 | try { |
836 | 7 | RETURN_IF_ERROR(target_type->get_serde()->read_column_from_arrow( |
837 | 7 | *columns[block_position.value()], array.get(), 0, num_rows, _ctz)); |
838 | 7 | } catch (const Exception& e) { |
839 | 0 | return Status::InternalError( |
840 | 0 | "Failed to convert ADBC Arrow column '{}' (file_column_id={}, arrow type={}) to " |
841 | 0 | "Doris block: {}", |
842 | 0 | column_name, file_column_id.value(), array->type()->ToString(), e.what()); |
843 | 0 | } |
844 | 7 | return Status::OK(); |
845 | 7 | } |
846 | | |
847 | 17 | Status AdbcFileReader::_build_col_name_to_file_id() { |
848 | 17 | _col_name_to_file_id.clear(); |
849 | 17 | _col_name_to_file_id.reserve(_file_slot_descs.size()); |
850 | 48 | for (size_t idx = 0; idx < _file_slot_descs.size(); ++idx) { |
851 | 31 | const auto* slot = _file_slot_descs[idx]; |
852 | 31 | DORIS_CHECK(slot != nullptr); |
853 | 31 | _col_name_to_file_id.emplace(slot->col_name(), LocalColumnId(cast_set<int32_t>(idx))); |
854 | 31 | } |
855 | 17 | return Status::OK(); |
856 | 17 | } |
857 | | |
858 | | AdbcReader::AdbcReader(AdbcStreamFactory stream_factory) |
859 | 1 | : _stream_factory(std::move(stream_factory)) {} |
860 | | |
861 | 1 | Status AdbcReader::init(TableReadOptions&& options) { |
862 | 1 | if (options.file_slot_descs == nullptr) { |
863 | 0 | return Status::InvalidArgument("ADBC reader requires file slot descriptors"); |
864 | 0 | } |
865 | 1 | return TableReader::init(std::move(options)); |
866 | 1 | } |
867 | | |
868 | 1 | Status AdbcReader::prepare_split(const SplitReadOptions& options) { |
869 | 1 | { |
870 | | // Keep protocol validation visible while avoiding overlap with TableReader's own scopes. |
871 | 1 | SCOPED_TIMER(_profile.total_timer); |
872 | 1 | SCOPED_TIMER(_profile.prepare_split_timer); |
873 | 1 | RETURN_IF_ERROR(validate_adbc_range(options.current_range)); |
874 | 1 | } |
875 | 1 | return TableReader::prepare_split(options); |
876 | 1 | } |
877 | | |
878 | 1 | Status AdbcReader::create_file_reader(std::unique_ptr<FileReader>* reader) { |
879 | 1 | DORIS_CHECK(reader != nullptr); |
880 | 1 | DORIS_CHECK(_file_slot_descs != nullptr); |
881 | 1 | *reader = std::make_unique<AdbcFileReader>(_system_properties, _current_task->data_file, |
882 | 1 | _io_ctx, _scanner_profile, _current_file_range_desc, |
883 | 1 | *_file_slot_descs, _stream_factory); |
884 | 1 | return Status::OK(); |
885 | 1 | } |
886 | | |
887 | | } // namespace doris::format::adbc |