be/src/format_v2/table/adbc_reader.cpp

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#include "format_v2/table/adbc_reader.h"

#include <arrow-adbc/adbc.h>
#include <arrow-adbc/adbc_driver_manager.h>
#include <arrow/array/array_base.h>
#include <arrow/c/bridge.h>
#include <arrow/record_batch.h>

#include <cstdint>
#include <map>
#include <memory>
#include <string>
#include <string_view>
#include <utility>
#include <vector>

#include "common/cast_set.h"
#include "common/check.h"
#include "core/assert_cast.h"
#include "core/block/block.h"
#include "core/data_type/data_type.h"
#include "core/data_type/data_type_array.h"
#include "core/data_type/data_type_map.h"
#include "core/data_type/data_type_nullable.h"
#include "core/data_type/data_type_struct.h"
#include "core/data_type/primitive_type.h"
#include "core/data_type_serde/data_type_serde.h"
#include "format/arrow/arrow_array_normalizer.h"
#include "format/parquet/arrow_memory_pool.h"
#include "format_v2/column_mapper.h"
#include "format_v2/materialized_reader_util.h"
#include "runtime/descriptors.h"
#include "runtime/exec_env.h"
#include "runtime/file_scan_profile.h"
#include "runtime/runtime_state.h"
#include "util/adbc_driver_registry.h"
#include "util/string_util.h"
#include "util/timezone_utils.h"
#include "util/url_coding.h"

namespace doris::format::adbc {
namespace {

// Keys of TTableFormatFileDesc.adbc_params. Kept next to the code that reads them so the FE-side
// producer and this consumer stay diffable.
constexpr const char* kParamDriverPath = "driver_path";
constexpr const char* kParamDriverEntrypoint = "driver_entrypoint";
constexpr const char* kParamUri = "uri";
constexpr const char* kParamUsername = "username";
constexpr const char* kParamPassword = "password";
constexpr const char* kParamQuerySql = "query_sql";
// Base64 of one opaque partition descriptor the driver produced on FE. Mutually exclusive with
// kParamQuerySql: a range either runs a statement here or reads one partition of a statement the
// source has already run.
constexpr const char* kParamPartitionDescriptor = "partition_descriptor";
// Anything under this prefix is an ADBC option name in full (the prefix is part of the option name,
// e.g. "adbc.connection.autocommit") and is handed to the driver untouched.
constexpr std::string_view kAdbcOptionPrefix = "adbc.";

const std::string* find_param(const std::map<std::string, std::string>& params,
                              const std::string& key) {
    const auto it = params.find(key);
    return it == params.end() ? nullptr : &it->second;
}

Status validate_adbc_range(const TFileRangeDesc& range) {
    if (!range.__isset.table_format_params ||
        range.table_format_params.table_format_type != "adbc") {
        return Status::InvalidArgument("ADBC reader requires the adbc table format");
    }
    if (!range.table_format_params.__isset.adbc_params) {
        return Status::InvalidArgument("ADBC reader requires adbc_params");
    }
    const auto& params = range.table_format_params.adbc_params;
    for (const auto* key : {kParamDriverPath, kParamUri}) {
        const auto* value = find_param(params, key);
        if (value == nullptr || value->empty()) {
            return Status::InvalidArgument("ADBC reader requires a non-empty '{}' parameter", key);
        }
    }
    const auto* query_sql = find_param(params, kParamQuerySql);
    const auto* partition = find_param(params, kParamPartitionDescriptor);
    const bool has_query = query_sql != nullptr && !query_sql->empty();
    const bool has_partition = partition != nullptr && !partition->empty();
    // Not a defensive nicety: reading a partition means the source has ALREADY run the statement, so
    // a range carrying both would let this reader run it a second time depending on which branch it
    // happened to take. FE refuses to build such a range; this refuses to act on one.
    if (has_query == has_partition) {
        return Status::InvalidArgument(
                "ADBC reader requires exactly one of '{}' and '{}', but the range carries {}",
                kParamQuerySql, kParamPartitionDescriptor, has_query ? "both" : "neither");
    }
    return Status::OK();
}

Status arrow_type_mismatch(const std::string& path, const arrow::DataType& arrow_type,
                           const DataTypePtr& doris_type) {
    return Status::InvalidArgument(
            "ADBC Arrow type mismatch at '{}': runtime type '{}' does not match cached Doris "
            "type '{}'",
            path, arrow_type.ToString(), doris_type->get_name());
}

int timestamp_scale(arrow::TimeUnit::type unit) {
    switch (unit) {
    case arrow::TimeUnit::SECOND:
        return 0;
    case arrow::TimeUnit::MILLI:
        return 3;
    case arrow::TimeUnit::MICRO:
    case arrow::TimeUnit::NANO:
        return 6;
    }
    return -1;
}

Status validate_arrow_type(const arrow::DataType& arrow_type, const DataTypePtr& doris_type,
                           const std::string& path) {
    DORIS_CHECK(doris_type != nullptr);
    const auto nested_doris_type = remove_nullable(doris_type);
    const auto primitive_type = nested_doris_type->get_primitive_type();
    const auto mismatch = [&]() { return arrow_type_mismatch(path, arrow_type, doris_type); };

    switch (arrow_type.id()) {
    case arrow::Type::BOOL:
        return primitive_type == TYPE_BOOLEAN ? Status::OK() : mismatch();
    case arrow::Type::INT8:
        return primitive_type == TYPE_TINYINT ? Status::OK() : mismatch();
    case arrow::Type::UINT8:
        return primitive_type == TYPE_SMALLINT ? Status::OK() : mismatch();
    case arrow::Type::INT16:
        return primitive_type == TYPE_SMALLINT ? Status::OK() : mismatch();
    case arrow::Type::UINT16:
        return primitive_type == TYPE_INT ? Status::OK() : mismatch();
    case arrow::Type::INT32:
        return primitive_type == TYPE_INT ? Status::OK() : mismatch();
    case arrow::Type::UINT32:
        return primitive_type == TYPE_BIGINT ? Status::OK() : mismatch();
    case arrow::Type::INT64:
        return primitive_type == TYPE_BIGINT ? Status::OK() : mismatch();
    case arrow::Type::UINT64:
        return primitive_type == TYPE_LARGEINT ? Status::OK() : mismatch();
    case arrow::Type::HALF_FLOAT:
    case arrow::Type::FLOAT:
        return primitive_type == TYPE_FLOAT ? Status::OK() : mismatch();
    case arrow::Type::DOUBLE:
        return primitive_type == TYPE_DOUBLE ? Status::OK() : mismatch();
    case arrow::Type::STRING:
    case arrow::Type::BINARY:
    case arrow::Type::FIXED_SIZE_BINARY:
        return primitive_type == TYPE_STRING ? Status::OK() : mismatch();
    case arrow::Type::DATE32:
        return primitive_type == TYPE_DATEV2 ? Status::OK() : mismatch();
    case arrow::Type::DATE64:
        return primitive_type == TYPE_DATETIMEV2 && nested_doris_type->get_scale() == 3
                       ? Status::OK()
                       : mismatch();
    case arrow::Type::TIMESTAMP: {
        const auto& timestamp = static_cast<const arrow::TimestampType&>(arrow_type);
        const bool zoned = !timestamp.timezone().empty();
        const auto expected = zoned ? TYPE_TIMESTAMPTZ : TYPE_DATETIMEV2;
        return primitive_type == expected &&
                               nested_doris_type->get_scale() == timestamp_scale(timestamp.unit())
                       ? Status::OK()
                       : mismatch();
    }
    case arrow::Type::DECIMAL128:
    case arrow::Type::DECIMAL256: {
        const auto& decimal = static_cast<const arrow::DecimalType&>(arrow_type);
        return is_decimal(primitive_type) &&
                               nested_doris_type->get_precision() == decimal.precision() &&
                               nested_doris_type->get_scale() == decimal.scale()
                       ? Status::OK()
                       : mismatch();
    }
    case arrow::Type::LIST:
    case arrow::Type::LARGE_LIST:
    case arrow::Type::FIXED_SIZE_LIST: {
        if (primitive_type != TYPE_ARRAY) {
            return mismatch();
        }
        const auto& list = static_cast<const arrow::BaseListType&>(arrow_type);
        const auto& array = assert_cast<const DataTypeArray&>(*nested_doris_type);
        return validate_arrow_type(*list.value_type(), array.get_nested_type(), path + ".element");
    }
    case arrow::Type::STRUCT: {
        if (primitive_type != TYPE_STRUCT) {
            return mismatch();
        }
        const auto& arrow_struct = static_cast<const arrow::StructType&>(arrow_type);
        const auto& doris_struct = assert_cast<const DataTypeStruct&>(*nested_doris_type);
        if (cast_set<size_t>(arrow_struct.num_fields()) != doris_struct.get_elements().size()) {
            return mismatch();
        }
        for (int field_idx = 0; field_idx < arrow_struct.num_fields(); ++field_idx) {
            const auto& arrow_field = arrow_struct.field(field_idx);
            const auto& doris_name = doris_struct.get_element_name(field_idx);
            if (to_lower(arrow_field->name()) != to_lower(doris_name)) {
                return Status::InvalidArgument(
                        "ADBC Arrow field mismatch at '{}': runtime field '{}' does not match "
                        "cached field '{}' at ordinal {}",
                        path, arrow_field->name(), doris_name, field_idx);
            }
            RETURN_IF_ERROR(validate_arrow_type(*arrow_field->type(),
                                                doris_struct.get_element(field_idx),
                                                path + "." + doris_name));
        }
        return Status::OK();
    }
    case arrow::Type::MAP: {
        if (primitive_type != TYPE_MAP) {
            return mismatch();
        }
        const auto& arrow_map = static_cast<const arrow::MapType&>(arrow_type);
        const auto& doris_map = assert_cast<const DataTypeMap&>(*nested_doris_type);
        RETURN_IF_ERROR(validate_arrow_type(*arrow_map.key_type(), doris_map.get_key_type(),
                                            path + ".key"));
        return validate_arrow_type(*arrow_map.item_type(), doris_map.get_value_type(),
                                   path + ".value");
    }
    default:
        return mismatch();
    }
}

// Drivers allocate the strings inside AdbcError, so every populated error has to be released.
class AdbcErrorGuard {
public:
    AdbcErrorGuard() = default;
    ~AdbcErrorGuard() { reset(); }
    AdbcErrorGuard(const AdbcErrorGuard&) = delete;
    AdbcErrorGuard& operator=(const AdbcErrorGuard&) = delete;

    AdbcError* get() { return &_error; }

    std::string take_message() {
        std::string message = _error.message != nullptr ? _error.message : "";
        reset();
        return message;
    }

    void reset() {
        if (_error.release != nullptr) {
            _error.release(&_error);
        }
        _error = ADBC_ERROR_INIT;
    }

private:
    AdbcError _error = ADBC_ERROR_INIT;
};

Status adbc_call_status(const char* what, AdbcStatusCode code, AdbcErrorGuard& error) {
    const std::string message = error.take_message();
    // AdbcStatusCode is a uint8_t, so spell out the name as well as the number.
    return Status::InternalError("ADBC: {} failed ({}, code {}): {}", what,
                                 AdbcStatusCodeMessage(code), static_cast<int>(code),
                                 message.empty() ? "driver reported no message" : message);
}

#define RETURN_IF_ADBC_ERROR(expr, what, error)                       \
    do {                                                              \
        const AdbcStatusCode adbc_call_code = (expr);                 \
        if (adbc_call_code != ADBC_STATUS_OK) {                       \
            return adbc_call_status((what), adbc_call_code, (error)); \
        }                                                             \
        (error).reset();                                              \
    } while (0)

// The production stream: one ADBC database/connection/statement per scan range.
//
// P0 keeps them un-pooled on purpose. Reusing databases across ranges is a throughput optimization
// that only pays off once multiple partitions run concurrently, and an unverifiable caching layer
// added now would only obscure the functional path.
class RealAdbcStream final : public AdbcStream {
public:
    explicit RealAdbcStream(const TFileRangeDesc& range) : _range(range) {}
    ~RealAdbcStream() override { static_cast<void>(close()); }

    Status open() {
        RETURN_IF_ERROR(validate_adbc_range(_range));
        const auto& params = _range.table_format_params.adbc_params;
        const std::string& driver_path = *find_param(params, kParamDriverPath);
        const std::string& uri = *find_param(params, kParamUri);
        const auto* entrypoint = find_param(params, kParamDriverEntrypoint);
        // validate_adbc_range has established that exactly one of these is present.
        const auto* partition = find_param(params, kParamPartitionDescriptor);

        RETURN_IF_ERROR(AdbcDriverRegistry::instance().get_or_load(
                driver_path, entrypoint != nullptr ? *entrypoint : std::string(), &_driver));
        DORIS_CHECK(_driver != nullptr);

        AdbcErrorGuard error;
        RETURN_IF_ADBC_ERROR(_driver->DatabaseNew(&_database, error.get()), "DatabaseNew", error);
        _database_created = true;
        RETURN_IF_ERROR(_set_database_options(params, uri, error));
        RETURN_IF_ADBC_ERROR(_driver->DatabaseInit(&_database, error.get()), "DatabaseInit", error);

        RETURN_IF_ADBC_ERROR(_driver->ConnectionNew(&_connection, error.get()), "ConnectionNew",
                             error);
        _connection_created = true;
        RETURN_IF_ADBC_ERROR(_driver->ConnectionInit(&_connection, &_database, error.get()),
                             "ConnectionInit", error);

        if (partition != nullptr && !partition->empty()) {
            RETURN_IF_ERROR(_read_partition(*partition, error));
        } else {
            RETURN_IF_ERROR(_execute_query(*find_param(params, kParamQuerySql), error));
        }

        // Before Arrow ever sees it: the driver's stream may not clear its release callback, and
        // Arrow aborts the process when that happens. Both branches need it -- the Flight SQL
        // driver's ReadPartition stream breaks the contract exactly like its ExecuteQuery one.
        enforce_stream_release_contract(&_c_stream);

        auto reader = arrow::ImportRecordBatchReader(&_c_stream);
        if (!reader.ok()) {
            return Status::InternalError("ADBC: failed to import the result stream: {}",
                                         reader.status().ToString());
        }
        // ImportRecordBatchReader moves the stream's contents; the reader owns it from here.
        _reader = reader.MoveValueUnsafe();
        return Status::OK();
    }

    Status next(std::shared_ptr<arrow::RecordBatch>* batch) override {
        DORIS_CHECK(batch != nullptr);
        if (_reader == nullptr) {
            return Status::InternalError("ADBC: result stream is not open");
        }
        std::shared_ptr<arrow::RecordBatch> next_batch;
        const auto status = _reader->ReadNext(&next_batch);
        if (!status.ok()) {
            return Status::InternalError("ADBC: failed to read the next batch: {}",
                                         status.ToString());
        }
        *batch = std::move(next_batch);
        return Status::OK();
    }

    Status close() override {
        Status result = Status::OK();
        // Release in reverse order of creation. The reader owns the imported stream, so it has to
        // go before the statement that produced it.
        _reader.reset();
        if (_c_stream.release != nullptr) {
            // Only reachable when the import itself failed; nothing else owns the stream then.
            _c_stream.release(&_c_stream);
            _c_stream = {};
        }
        AdbcErrorGuard error;
        if (_statement_created) {
            const auto code = _driver->StatementRelease(&_statement, error.get());
            if (code != ADBC_STATUS_OK && result.ok()) {
                result = adbc_call_status("StatementRelease", code, error);
            }
            error.reset();
            _statement_created = false;
        }
        if (_connection_created) {
            const auto code = _driver->ConnectionRelease(&_connection, error.get());
            if (code != ADBC_STATUS_OK && result.ok()) {
                result = adbc_call_status("ConnectionRelease", code, error);
            }
            error.reset();
            _connection_created = false;
        }
        if (_database_created) {
            const auto code = _driver->DatabaseRelease(&_database, error.get());
            if (code != ADBC_STATUS_OK && result.ok()) {
                result = adbc_call_status("DatabaseRelease", code, error);
            }
            error.reset();
            _database_created = false;
        }
        // _driver itself is owned by AdbcDriverRegistry and is never released.
        return result;
    }

private:
    // Runs the statement FE generated. One statement per range, so this range is the whole query.
    Status _execute_query(const std::string& query_sql, AdbcErrorGuard& error) {
        RETURN_IF_ADBC_ERROR(_driver->StatementNew(&_connection, &_statement, error.get()),
                             "StatementNew", error);
        _statement_created = true;
        RETURN_IF_ADBC_ERROR(
                _driver->StatementSetSqlQuery(&_statement, query_sql.c_str(), error.get()),
                "StatementSetSqlQuery", error);
        int64_t rows_affected = -1;
        RETURN_IF_ADBC_ERROR(_driver->StatementExecuteQuery(&_statement, &_c_stream, &rows_affected,
                                                            error.get()),
                             "StatementExecuteQuery", error);
        return Status::OK();
    }

    // Reads one partition of a query FE already had the source execute. No statement is created:
    // ADBC reads a partition off a connection, and the whole point is that this can happen on a
    // different machine from the one that planned it.
    Status _read_partition(const std::string& base64_descriptor, AdbcErrorGuard& error) {
        std::string descriptor;
        if (!base64_decode(base64_descriptor, &descriptor)) {
            return Status::InvalidArgument("ADBC: the '{}' parameter is not valid base64",
                                           kParamPartitionDescriptor);
        }
        RETURN_IF_ADBC_ERROR(
                _driver->ConnectionReadPartition(
                        &_connection, reinterpret_cast<const uint8_t*>(descriptor.data()),
                        descriptor.size(), &_c_stream, error.get()),
                "ConnectionReadPartition", error);
        return Status::OK();
    }

    Status _set_database_options(const std::map<std::string, std::string>& params,
                                 const std::string& uri, AdbcErrorGuard& error) {
        RETURN_IF_ADBC_ERROR(
                _driver->DatabaseSetOption(&_database, ADBC_OPTION_URI, uri.c_str(), error.get()),
                "DatabaseSetOption(uri)", error);
        for (const auto* key : {kParamUsername, kParamPassword}) {
            const auto* value = find_param(params, key);
            if (value == nullptr || value->empty()) {
                continue;
            }
            RETURN_IF_ADBC_ERROR(
                    _driver->DatabaseSetOption(&_database, key, value->c_str(), error.get()),
                    "DatabaseSetOption(credentials)", error);
        }
        for (const auto& [key, value] : params) {
            if (!key.starts_with(kAdbcOptionPrefix)) {
                continue;
            }
            RETURN_IF_ADBC_ERROR(
                    _driver->DatabaseSetOption(&_database, key.c_str(), value.c_str(), error.get()),
                    "DatabaseSetOption(passthrough)", error);
        }
        return Status::OK();
    }

    const TFileRangeDesc _range;
    const AdbcDriver* _driver = nullptr;
    AdbcDatabase _database {};
    AdbcConnection _connection {};
    AdbcStatement _statement {};
    ArrowArrayStream _c_stream {};
    std::shared_ptr<arrow::RecordBatchReader> _reader;
    bool _database_created = false;
    bool _connection_created = false;
    bool _statement_created = false;
};

Status create_real_adbc_stream(const TFileRangeDesc& range, std::unique_ptr<AdbcStream>* out) {
    DORIS_CHECK(out != nullptr);
    auto stream = std::make_unique<RealAdbcStream>(range);
    RETURN_IF_ERROR(stream->open());
    *out = std::move(stream);
    return Status::OK();
}

ColumnDefinition adbc_child_definition(const std::string& name, DataTypePtr type, int32_t local_id);

// Mirrors synthesize_remote_doris_children in remote_doris_reader.cpp. Both readers expose table
// slots as file columns, so complex columns still need structural children for TableColumnMapper.
// Kept separate rather than shared to avoid reshaping the already-shipped remote_doris reader.
std::vector<ColumnDefinition> synthesize_adbc_children(const DataTypePtr& type) {
    std::vector<ColumnDefinition> children;
    DORIS_CHECK(type != nullptr);
    const auto nested_type = remove_nullable(type);
    switch (nested_type->get_primitive_type()) {
    case TYPE_ARRAY: {
        const auto* array_type = assert_cast<const DataTypeArray*>(nested_type.get());
        children.push_back(adbc_child_definition("element", array_type->get_nested_type(), 0));
        break;
    }
    case TYPE_MAP: {
        const auto* map_type = assert_cast<const DataTypeMap*>(nested_type.get());
        children.push_back(adbc_child_definition("key", map_type->get_key_type(), 0));
        children.push_back(adbc_child_definition("value", map_type->get_value_type(), 1));
        break;
    }
    case TYPE_STRUCT: {
        const auto* struct_type = assert_cast<const DataTypeStruct*>(nested_type.get());
        children.reserve(struct_type->get_elements().size());
        for (size_t idx = 0; idx < struct_type->get_elements().size(); ++idx) {
            children.push_back(adbc_child_definition(struct_type->get_element_name(idx),
                                                     struct_type->get_element(idx),
                                                     cast_set<int32_t>(idx)));
        }
        break;
    }
    default:
        break;
    }
    return children;
}

ColumnDefinition adbc_child_definition(const std::string& name, DataTypePtr type,
                                       int32_t local_id) {
    ColumnDefinition child;
    child.identifier = Field::create_field<TYPE_STRING>(name);
    child.local_id = local_id;
    child.name = name;
    child.type = std::move(type);
    child.children = synthesize_adbc_children(child.type);
    return child;
}

// A stream that forwards everything to the driver's and, on release, does the one thing some
// drivers forget: clear its own release callback. Heap-allocated because Arrow keeps only the
// ArrowArrayStream it was handed, and the delegate has to outlive this function.
struct DelegatingStream {
    ArrowArrayStream inner;
};

int delegating_get_schema(ArrowArrayStream* self, ArrowSchema* out) {
    auto& inner = static_cast<DelegatingStream*>(self->private_data)->inner;
    return inner.get_schema(&inner, out);
}

int delegating_get_next(ArrowArrayStream* self, ArrowArray* out) {
    auto& inner = static_cast<DelegatingStream*>(self->private_data)->inner;
    return inner.get_next(&inner, out);
}

const char* delegating_get_last_error(ArrowArrayStream* self) {
    auto& inner = static_cast<DelegatingStream*>(self->private_data)->inner;
    return inner.get_last_error != nullptr ? inner.get_last_error(&inner) : nullptr;
}

void delegating_release(ArrowArrayStream* self) {
    auto* delegate = static_cast<DelegatingStream*>(self->private_data);
    if (delegate->inner.release != nullptr) {
        delegate->inner.release(&delegate->inner);
    }
    delete delegate;
    self->private_data = nullptr;
    // What the driver failed to do, and what Arrow aborts the process over.
    self->release = nullptr;
}

} // namespace

void enforce_stream_release_contract(ArrowArrayStream* stream) {
    DORIS_CHECK(stream != nullptr);
    if (stream->release == nullptr) {
        // Already released; nothing to delegate to, and wrapping it would hand Arrow a stream
        // whose callbacks dereference a released delegate.
        return;
    }
    auto* delegate = new DelegatingStream {.inner = *stream};
    *stream = ArrowArrayStream {.get_schema = delegating_get_schema,
                                .get_next = delegating_get_next,
                                .get_last_error = delegating_get_last_error,
                                .release = delegating_release,
                                .private_data = delegate};
}

AdbcFileReader::AdbcFileReader(std::shared_ptr<io::FileSystemProperties>& system_properties,
                               std::unique_ptr<io::FileDescription>& file_description,
                               std::shared_ptr<io::IOContext> io_ctx, RuntimeProfile* profile,
                               const TFileRangeDesc& range,
                               const std::vector<SlotDescriptor*>& file_slot_descs,
                               AdbcStreamFactory stream_factory)
        : FileReader(system_properties, file_description, std::move(io_ctx), profile),
          _range(range),
          _file_slot_descs(file_slot_descs),
          _stream_factory(std::move(stream_factory)) {
    TimezoneUtils::find_cctz_time_zone(TimezoneUtils::default_time_zone, _ctz);
}

AdbcFileReader::~AdbcFileReader() {
    static_cast<void>(close());
}

void AdbcFileReader::_init_profile() {
    if (_profile == nullptr) {
        return;
    }
    const auto hierarchy = file_scan_profile::ensure_hierarchy(_profile);
    _io_time = hierarchy.io;
    static const char* adbc_profile = "AdbcFileReader";
    _total_time =
            ADD_CHILD_TIMER_WITH_LEVEL(_profile, adbc_profile, file_scan_profile::FILE_READER, 1);
    _open_stream_time = ADD_CHILD_TIMER_WITH_LEVEL(_profile, "AdbcOpenStreamTime", adbc_profile, 1);
    _next_batch_time = ADD_CHILD_TIMER_WITH_LEVEL(_profile, "AdbcNextBatchTime", adbc_profile, 1);
    _normalize_time = ADD_CHILD_TIMER_WITH_LEVEL(_profile, "AdbcNormalizeTime", adbc_profile, 1);
    _materialize_time =
            ADD_CHILD_TIMER_WITH_LEVEL(_profile, "AdbcMaterializeTime", adbc_profile, 1);
    _filter_time = ADD_CHILD_TIMER_WITH_LEVEL(_profile, "AdbcFilterTime", adbc_profile, 1);
}

Status AdbcFileReader::init(RuntimeState* state) {
    _init_profile();
    SCOPED_TIMER(_total_time);
    _runtime_state = state;
    RETURN_IF_ERROR(validate_adbc_range(_range));
    RETURN_IF_ERROR(_build_col_name_to_file_id());
    _eof = false;
    return Status::OK();
}

Status AdbcFileReader::get_schema(std::vector<ColumnDefinition>* file_schema) const {
    SCOPED_TIMER(_total_time);
    DORIS_CHECK(file_schema != nullptr);
    file_schema->clear();
    file_schema->reserve(_file_slot_descs.size());
    for (size_t idx = 0; idx < _file_slot_descs.size(); ++idx) {
        const auto* slot = _file_slot_descs[idx];
        DORIS_CHECK(slot != nullptr);
        file_schema->push_back({
                .identifier = Field::create_field<TYPE_INT>(cast_set<int32_t>(idx)),
                .local_id = cast_set<int32_t>(idx),
                .name = slot->col_name(),
                .type = slot->type(),
                .children = synthesize_adbc_children(slot->type()),
        });
    }
    return Status::OK();
}

std::unique_ptr<TableColumnMapper> AdbcFileReader::create_column_mapper(
        TableColumnMapperOptions options) const {
    // ADBC streams return complete Arrow roots, so TableReader must own every nested projection and
    // reorder after the full source shape has been materialized.
    return std::make_unique<MaterializedColumnMapper>(std::move(options));
}

Status AdbcFileReader::open(std::shared_ptr<FileScanRequest> request) {
    SCOPED_TIMER(_total_time);
    SCOPED_TIMER(_open_stream_time);
    RETURN_IF_ERROR(FileReader::open(std::move(request)));
    RETURN_IF_ERROR(_open_stream());
    _eof = false;
    return Status::OK();
}

Status AdbcFileReader::get_block(Block* file_block, size_t* rows, bool* eof) {
    SCOPED_TIMER(_total_time);
    DORIS_CHECK(file_block != nullptr);
    DORIS_CHECK(rows != nullptr);
    DORIS_CHECK(eof != nullptr);
    if (_stream == nullptr) {
        return Status::InternalError("ADBC reader is not open");
    }
    if (_io_ctx != nullptr && _io_ctx->should_stop) {
        // Observe cancellation before entering a potentially blocking driver read.
        RETURN_IF_ERROR(close());
        *rows = 0;
        *eof = true;
        return Status::OK();
    }

    *rows = 0;
    *eof = false;
    std::shared_ptr<arrow::RecordBatch> batch;
    {
        SCOPED_TIMER(_io_time);
        SCOPED_TIMER(_next_batch_time);
        RETURN_IF_ERROR(_stream->next(&batch));
    }
    if (batch == nullptr) {
        *eof = true;
        _eof = true;
        return Status::OK();
    }

    {
        SCOPED_TIMER(_materialize_time);
        RETURN_IF_ERROR(_materialize_record_batch(*batch, file_block, rows));
    }
    _record_scan_rows(cast_set<int64_t>(*rows));
    {
        SCOPED_TIMER(_filter_time);
        RETURN_IF_ERROR(
                apply_materialized_reader_filters(_request.get(), _io_ctx.get(), file_block, rows));
    }
    return Status::OK();
}

Status AdbcFileReader::close() {
    SCOPED_TIMER(_total_time);
    if (_stream != nullptr) {
        RETURN_IF_ERROR(_stream->close());
        _stream.reset();
    }
    _request.reset();
    _eof = true;
    return Status::OK();
}

Status AdbcFileReader::_open_stream() {
    DORIS_CHECK(_stream == nullptr);
    if (_stream_factory) {
        RETURN_IF_ERROR(_stream_factory(_range, &_stream));
    } else {
        RETURN_IF_ERROR(create_real_adbc_stream(_range, &_stream));
    }
    DORIS_CHECK(_stream != nullptr);
    return Status::OK();
}

Status AdbcFileReader::_materialize_record_batch(const arrow::RecordBatch& batch, Block* file_block,
                                                 size_t* rows) const {
    DORIS_CHECK(file_block != nullptr);
    DORIS_CHECK(rows != nullptr);
    if (_request == nullptr) {
        return Status::InternalError("ADBC reader is not open");
    }

    if (_col_name_to_file_id.empty()) {
        // A pushed-down COUNT(*) projects no columns at all: the scan wants rows counted, no values.
        // Counting here rather than falling into the loop below is not an optimization -- every column
        // the source returns is unrequested by definition in this state, so the loop's unknown-column
        // check would reject the first one and fail a query that asked for nothing but a number.
        // FE sends a one-constant-column statement for this case, so the batch is narrow.
        //
        // Only the empty case is special-cased. An unrequested column arriving alongside requested ones
        // still fails: that means FE and this reader disagree about the projection, and it is the one
        // signal that the disagreement exists.
        *rows = cast_set<size_t>(batch.num_rows());
        return Status::OK();
    }

    ArrowMemoryPool<> local_arrow_pool;
    arrow::MemoryPool* arrow_pool = ExecEnv::GetInstance()->arrow_memory_pool();
    if (arrow_pool == nullptr) {
        // Embedded and unit-test runtimes may omit ExecEnv memory initialization; keep conversions
        // on Doris' tracked allocator instead of falling back to Arrow's untracked default pool.
        arrow_pool = &local_arrow_pool;
    }
    std::vector<bool> materialized_columns(file_block->columns(), false);
    for (int arrow_idx = 0; arrow_idx < batch.num_columns(); ++arrow_idx) {
        const std::string& column_name = batch.schema()->field(arrow_idx)->name();
        const auto file_id_it = _col_name_to_file_id.find(column_name);
        if (file_id_it == _col_name_to_file_id.end()) {
            return Status::InternalError("ADBC source returned unknown column {}", column_name);
        }
        const auto block_position_it = _request->local_positions.find(file_id_it->second);
        if (block_position_it == _request->local_positions.end()) {
            continue;
        }
        std::shared_ptr<arrow::Array> array;
        {
            SCOPED_TIMER(_normalize_time);
            RETURN_IF_ERROR(normalize_arrow_array(batch.column(arrow_idx), arrow_pool, &array));
        }
        RETURN_IF_ERROR(_materialize_arrow_column(column_name, array, batch.num_rows(),
                                                  file_id_it->second, block_position_it->second,
                                                  file_block));
        materialized_columns[block_position_it->second.value()] = true;
    }

    for (const auto& [file_column_id, block_position] : _request->local_positions) {
        if (block_position.value() >= materialized_columns.size()) {
            return Status::InternalError(
                    "ADBC requested block position {} out of range, block columns {}",
                    block_position.value(), materialized_columns.size());
        }
        if (!materialized_columns[block_position.value()]) {
            return Status::InternalError("ADBC source did not return requested file column id {}",
                                         file_column_id.value());
        }
    }

    *rows = cast_set<size_t>(batch.num_rows());
    return Status::OK();
}

Status AdbcFileReader::_materialize_arrow_column(const std::string& column_name,
                                                 const std::shared_ptr<arrow::Array>& array,
                                                 int64_t num_rows, LocalColumnId file_column_id,
                                                 const LocalIndex& block_position,
                                                 Block* file_block) const {
    DORIS_CHECK(file_block != nullptr);
    DORIS_CHECK(array != nullptr);
    if (block_position.value() >= file_block->columns()) {
        return Status::InternalError("ADBC block position {} out of range, block columns {}",
                                     block_position.value(), file_block->columns());
    }
    const auto& target_type = file_block->get_by_position(block_position.value()).type;

    // The cached FE target remains authoritative when a source schema changes mid-cache-window;
    // otherwise SerDes without a null map can silently turn source nulls into default values.
    if (array->null_count() > 0 && !target_type->is_nullable()) {
        return Status::InternalError(
                "ADBC Arrow column '{}' contains {} null rows for non-nullable Doris type {}",
                column_name, array->null_count(), target_type->get_name());
    }

    // An all-null column arrives with a type that says nothing about the column.
    //
    // A source that infers Arrow types from the VALUES it returns -- rather than from the declared
    // column type -- has nothing to infer from when every value in the result is null, and picks
    // whatever its default is. Measured against the SQLite driver: the same TEXT column comes back
    // as utf8 for `SELECT id, name FROM t1` and as int64 for
    // `SELECT id, name FROM t1 WHERE name IS NULL`, purely because the filter left only nulls.
    // Handing that to the serde fails with "Unsupported arrow type for string column: 9", and no
    // amount of care on the FE side avoids it: FE cannot know in advance which rows will survive.
    //
    // N nulls are what this array means whatever type it claims, so materialize them directly. The
    // check is narrow on purpose -- a column with even one non-null value keeps its real type and
    // still fails loudly on a genuine mismatch, which is the signal that FE and the source disagree
    // about the schema rather than about one result set.
    //
    // Only for a nullable target: substituting defaults into a NOT NULL column would turn a source
    // that wrongly sent nulls into silently wrong data, so that keeps failing in the serde.
    if (array->null_count() == array->length() && target_type->is_nullable()) {
        auto columns_guard = file_block->mutate_columns_scoped();
        auto& columns = columns_guard.mutable_columns();
        columns[block_position.value()]->insert_many_defaults(cast_set<size_t>(num_rows));
        return Status::OK();
    }

    // Validate the whole logical shape before acquiring mutable block columns. Struct SerDes map
    // children by ordinal, so a stale cache entry can otherwise produce plausible but wrong values.
    RETURN_IF_ERROR(validate_arrow_type(*array->type(), target_type, column_name));

    auto columns_guard = file_block->mutate_columns_scoped();
    auto& columns = columns_guard.mutable_columns();
    try {
        RETURN_IF_ERROR(target_type->get_serde()->read_column_from_arrow(
                *columns[block_position.value()], array.get(), 0, num_rows, _ctz));
    } catch (const Exception& e) {
        return Status::InternalError(
                "Failed to convert ADBC Arrow column '{}' (file_column_id={}, arrow type={}) to "
                "Doris block: {}",
                column_name, file_column_id.value(), array->type()->ToString(), e.what());
    }
    return Status::OK();
}

Status AdbcFileReader::_build_col_name_to_file_id() {
    _col_name_to_file_id.clear();
    _col_name_to_file_id.reserve(_file_slot_descs.size());
    for (size_t idx = 0; idx < _file_slot_descs.size(); ++idx) {
        const auto* slot = _file_slot_descs[idx];
        DORIS_CHECK(slot != nullptr);
        _col_name_to_file_id.emplace(slot->col_name(), LocalColumnId(cast_set<int32_t>(idx)));
    }
    return Status::OK();
}

AdbcReader::AdbcReader(AdbcStreamFactory stream_factory)
        : _stream_factory(std::move(stream_factory)) {}

Status AdbcReader::init(TableReadOptions&& options) {
    if (options.file_slot_descs == nullptr) {
        return Status::InvalidArgument("ADBC reader requires file slot descriptors");
    }
    return TableReader::init(std::move(options));
}

Status AdbcReader::prepare_split(const SplitReadOptions& options) {
    {
        // Keep protocol validation visible while avoiding overlap with TableReader's own scopes.
        SCOPED_TIMER(_profile.total_timer);
        SCOPED_TIMER(_profile.prepare_split_timer);
        RETURN_IF_ERROR(validate_adbc_range(options.current_range));
    }
    return TableReader::prepare_split(options);
}

Status AdbcReader::create_file_reader(std::unique_ptr<FileReader>* reader) {
    DORIS_CHECK(reader != nullptr);
    DORIS_CHECK(_file_slot_descs != nullptr);
    *reader = std::make_unique<AdbcFileReader>(_system_properties, _current_task->data_file,
                                               _io_ctx, _scanner_profile, _current_file_range_desc,
                                               *_file_slot_descs, _stream_factory);
    return Status::OK();
}

} // namespace doris::format::adbc