Coverage Report

Created: 2026-08-21 22:52

next uncovered line (L), next uncovered region (R), next uncovered branch (B)
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