be/src/format_v2/wal/wal_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/wal/wal_reader.h"

#include <absl/strings/numbers.h>
#include <absl/strings/str_split.h>

#include <ranges>
#include <unordered_set>
#include <utility>

#include "agent/be_exec_version_manager.h"
#include "common/cast_set.h"
#include "core/block/block.h"
#include "core/data_type/data_type_factory.hpp"
#include "core/data_type/data_type_nullable.h"
#include "format_v2/column_mapper.h"
#include "format_v2/materialized_reader_util.h"
#include "load/group_commit/wal/wal_file_reader.h"
#include "load/group_commit/wal/wal_manager.h"
#include "runtime/exec_env.h"
#include "runtime/runtime_state.h"

namespace doris::format::wal {
namespace {

class WalColumnMapper final : public TableColumnMapper {
public:
    using TableColumnMapper::TableColumnMapper;

    Status create_mapping(const std::vector<ColumnDefinition>& projected_columns,
                          const std::map<std::string, Field>& partition_values,
                          const std::vector<ColumnDefinition>& file_schema) override {
        for (const auto& projected : projected_columns) {
            if (!projected.has_identifier_field_id()) {
                return Status::InternalError("WAL projected column {} has no unique id",
                                             projected.name);
            }
            const auto found = std::ranges::find_if(file_schema, [&](const auto& file_column) {
                return file_column.has_identifier_field_id() &&
                       file_column.get_identifier_field_id() == projected.get_identifier_field_id();
            });
            if (found == file_schema.end()) {
                return Status::InternalError("WAL does not contain column unique id {} ({})",
                                             projected.get_identifier_field_id(), projected.name);
            }
        }
        return TableColumnMapper::create_mapping(projected_columns, partition_values, file_schema);
    }

protected:
    bool enable_lazy_materialization() const override { return false; }
    bool force_full_complex_scan_projection() const override { return true; }
};

ColumnDefinition build_wal_column_definition(const PColumnMeta& meta, int32_t local_id) {
    ColumnDefinition field;
    field.local_id = local_id;
    field.name = meta.name();
    field.type = make_nullable(DataTypeFactory::instance().create_data_type(meta));
    field.children.reserve(meta.children_size());
    for (int child_idx = 0; child_idx < meta.children_size(); ++child_idx) {
        field.children.push_back(build_wal_column_definition(meta.children(child_idx), child_idx));
    }
    return field;
}

} // namespace

Status parse_wal_column_ids(const std::string& encoded, std::vector<int32_t>* column_ids) {
    DORIS_CHECK(column_ids != nullptr);
    column_ids->clear();
    if (encoded.empty()) {
        return Status::Corruption("WAL header contains no column ids");
    }

    std::unordered_set<int32_t> seen;
    for (const absl::string_view token : absl::StrSplit(encoded, ',')) {
        int32_t column_id = 0;
        if (token.empty() || !absl::SimpleAtoi(token, &column_id)) {
            return Status::Corruption("invalid WAL column id '{}'", std::string(token));
        }
        if (!seen.emplace(column_id).second) {
            return Status::Corruption("duplicate WAL column id {}", column_id);
        }
        column_ids->push_back(column_id);
    }
    return Status::OK();
}

WalReader::WalReader(std::shared_ptr<io::FileSystemProperties>& system_properties,
                     std::unique_ptr<io::FileDescription>& file_description,
                     std::shared_ptr<io::IOContext> io_ctx, RuntimeProfile* profile,
                     const std::vector<ColumnDefinition>& projected_columns)
        : FileReader(system_properties, file_description, std::move(io_ctx), profile),
          _projected_columns(projected_columns) {}

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

Status WalReader::init(RuntimeState* state) {
    if (state == nullptr || state->exec_env() == nullptr ||
        state->exec_env()->wal_mgr() == nullptr) {
        return Status::InvalidArgument("WAL v2 reader requires a runtime WAL manager");
    }
    RETURN_IF_ERROR(state->exec_env()->wal_mgr()->get_wal_path(state->wal_id(), _wal_path));
    _wal_reader = std::make_shared<doris::WalFileReader>(_wal_path);
    RETURN_IF_ERROR(_wal_reader->init());

    std::string encoded_column_ids;
    RETURN_IF_ERROR(_wal_reader->read_header(_version, encoded_column_ids));
    RETURN_IF_ERROR(parse_wal_column_ids(encoded_column_ids, &_column_ids));
    _reader_eof = false;
    _eof = false;
    return Status::OK();
}

Status WalReader::get_schema(std::vector<ColumnDefinition>* file_schema) const {
    if (file_schema == nullptr) {
        return Status::InvalidArgument("WAL v2 file_schema is null");
    }
    RETURN_IF_ERROR(_ensure_schema_loaded());
    *file_schema = _file_schema;
    return Status::OK();
}

std::unique_ptr<TableColumnMapper> WalReader::create_column_mapper(
        TableColumnMapperOptions options) const {
    return std::make_unique<WalColumnMapper>(std::move(options));
}

Status WalReader::open(std::shared_ptr<FileScanRequest> request) {
    RETURN_IF_ERROR(FileReader::open(std::move(request)));
    _first_block_consumed = false;
    _eof = false;
    return Status::OK();
}

Status WalReader::get_block(Block* file_block, size_t* rows, bool* eof) {
    DORIS_CHECK(file_block != nullptr);
    DORIS_CHECK(rows != nullptr);
    DORIS_CHECK(eof != nullptr);
    if (_request == nullptr) {
        return Status::InternalError("WAL v2 reader is not open");
    }

    *rows = 0;
    *eof = false;
    if (_reader_eof) {
        *eof = true;
        _eof = true;
        return Status::OK();
    }

    PBlock pblock;
    if (_first_block_loaded && !_first_block_consumed) {
        // Schema discovery owns the first payload temporarily; transfer it into the read path so
        // the reader neither retains a second PBlock nor forces protobuf to clone its buffers.
        pblock.Swap(&_first_block);
        _first_block_consumed = true;
        _first_block_loaded = false;
    } else {
        auto status = _wal_reader->read_block(pblock);
        if (status.is<ErrorCode::END_OF_FILE>()) {
            _reader_eof = true;
            *eof = true;
            _eof = true;
            return Status::OK();
        }
        RETURN_IF_ERROR(status);
    }
    RETURN_IF_ERROR(_validate_block_version(pblock));

    Block source_block;
    size_t uncompressed_size = 0;
    int64_t decompress_time = 0;
    RETURN_IF_ERROR(source_block.deserialize(pblock, &uncompressed_size, &decompress_time));
    if (source_block.columns() != _column_ids.size()) {
        return Status::Corruption("WAL block has {} columns but header declares {}",
                                  source_block.columns(), _column_ids.size());
    }
    RETURN_IF_ERROR(_materialize_requested_columns(&source_block, file_block));
    *rows = file_block->rows();
    _record_scan_rows(cast_set<int64_t>(*rows));
    RETURN_IF_ERROR(
            apply_materialized_reader_filters(_request.get(), _io_ctx.get(), file_block, rows));
    return Status::OK();
}

Status WalReader::close() {
    _request.reset();
    _reader_eof = true;
    _eof = true;
    if (_wal_reader == nullptr) {
        return Status::OK();
    }
    auto status = _wal_reader->finalize();
    if (status.ok()) {
        _wal_reader.reset();
    }
    return status;
}

Status WalReader::_ensure_schema_loaded() const {
    if (_schema_inited) {
        return Status::OK();
    }

    auto status = _wal_reader->read_block(_first_block);
    if (status.is<ErrorCode::END_OF_FILE>()) {
        // An empty WAL still has a complete unique-id header. Use only matching projected types;
        // there is no data block from which unprojected physical types could be inferred.
        return _init_schema_from_block(nullptr);
    }
    RETURN_IF_ERROR(status);
    RETURN_IF_ERROR(_validate_block_version(_first_block));
    _first_block_loaded = true;
    return _init_schema_from_block(&_first_block);
}

Status WalReader::_validate_block_version(const PBlock& pblock) const {
    const int version = pblock.has_be_exec_version() ? pblock.be_exec_version() : 0;
    if (!BeExecVersionManager::check_be_exec_version(version)) {
        return Status::DataQualityError("unsupported BE execution version {} in WAL", version);
    }
    return Status::OK();
}

Status WalReader::_init_schema_from_block(const PBlock* pblock) const {
    if (pblock != nullptr && cast_set<size_t>(pblock->column_metas_size()) != _column_ids.size()) {
        return Status::Corruption("WAL block schema has {} columns but header declares {}",
                                  pblock->column_metas_size(), _column_ids.size());
    }

    _file_schema.clear();
    for (size_t idx = 0; idx < _column_ids.size(); ++idx) {
        ColumnDefinition field;
        field.identifier = Field::create_field<TYPE_INT>(_column_ids[idx]);
        field.local_id = cast_set<int32_t>(idx);
        if (pblock != nullptr) {
            const auto& meta = pblock->column_metas(cast_set<int>(idx));
            // WAL metadata is the file-local schema; preserve its complete nested shape so the
            // mapper can validate ARRAY/MAP/STRUCT projections instead of seeing an empty shell.
            field = build_wal_column_definition(meta, cast_set<int32_t>(idx));
            field.identifier = Field::create_field<TYPE_INT>(_column_ids[idx]);
        } else {
            const auto projected =
                    std::ranges::find_if(_projected_columns, [&](const auto& candidate) {
                        return candidate.has_identifier_field_id() &&
                               candidate.get_identifier_field_id() == _column_ids[idx];
                    });
            if (projected == _projected_columns.end()) {
                continue;
            }
            field.name = projected->name;
            field.type = projected->type;
        }
        _file_schema.push_back(std::move(field));
    }
    _schema_inited = true;
    return Status::OK();
}

Status WalReader::_materialize_requested_columns(Block* source_block, Block* file_block) const {
    DORIS_CHECK(source_block != nullptr);
    for (const auto& [file_column_id, block_position] : _request->local_positions) {
        const auto source_idx = file_column_id.value();
        if (source_idx < 0 || cast_set<size_t>(source_idx) >= source_block->columns()) {
            return Status::Corruption("WAL request refers to invalid local column {}", source_idx);
        }
        if (block_position.value() >= file_block->columns()) {
            return Status::InternalError("WAL request has invalid block position {}",
                                         block_position.value());
        }
        const auto& target = file_block->get_by_position(block_position.value());
        // Deserialized WAL columns have a single owner. Move that ownership into the output block
        // before mutate() so wide payloads do not trigger copy-on-write deep clones.
        auto column = std::move(source_block->get_by_position(source_idx).column);
        column = make_column_nullable_if_needed(std::move(column), target.type);
        file_block->replace_by_position(block_position.value(), IColumn::mutate(std::move(column)));
    }
    return Status::OK();
}

} // namespace doris::format::wal