be/src/format_v2/native/native_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/native/native_reader.h"

#include <cstring>
#include <utility>

#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/native/native_format.h"
#include "format_v2/column_mapper.h"
#include "format_v2/materialized_reader_util.h"
#include "io/file_factory.h"
#include "io/fs/tracing_file_reader.h"
#include "runtime/file_scan_profile.h"
#include "runtime/runtime_state.h"
#include "util/slice.h"

namespace doris::format::native {
namespace {

Status parse_native_pblock(const std::string& buffer, const std::string& path, PBlock* pblock) {
    DORIS_CHECK(pblock != nullptr);
    if (!pblock->ParseFromArray(buffer.data(), cast_set<int>(buffer.size()))) {
        return Status::InternalError("Failed to parse native PBlock from file {}", path);
    }
    return Status::OK();
}

} // namespace

NativeReader::NativeReader(std::shared_ptr<io::FileSystemProperties>& system_properties,
                           std::unique_ptr<io::FileDescription>& file_description,
                           std::shared_ptr<io::IOContext> io_ctx, RuntimeProfile* profile)
        : FileReader(system_properties, file_description, std::move(io_ctx), profile) {}

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

void NativeReader::_init_profile() {
    if (_profile == nullptr) {
        return;
    }
    file_scan_profile::ensure_hierarchy(_profile);
    static const char* native_profile = "NativeReader";
    _total_time =
            ADD_CHILD_TIMER_WITH_LEVEL(_profile, native_profile, file_scan_profile::FILE_READER, 1);
    _read_block_time =
            ADD_CHILD_TIMER_WITH_LEVEL(_profile, "NativeReadBlockTime", native_profile, 1);
    _deserialize_time =
            ADD_CHILD_TIMER_WITH_LEVEL(_profile, "NativeDeserializeTime", native_profile, 1);
    _materialize_time =
            ADD_CHILD_TIMER_WITH_LEVEL(_profile, "NativeMaterializeTime", native_profile, 1);
    _filter_time = ADD_CHILD_TIMER_WITH_LEVEL(_profile, "NativeFilterTime", native_profile, 1);
}

Status NativeReader::init(RuntimeState* state) {
    _init_profile();
    SCOPED_TIMER(_total_time);
    _runtime_state = state;
    if (_file_description == nullptr) {
        return Status::InvalidArgument("Native v2 reader requires file description");
    }
    RETURN_IF_ERROR(FileReader::init(state));
    RETURN_IF_ERROR(_validate_and_consume_header());
    return Status::OK();
}

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

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

Status NativeReader::open(std::shared_ptr<FileScanRequest> request) {
    SCOPED_TIMER(_total_time);
    RETURN_IF_ERROR(FileReader::open(std::move(request)));
    DORIS_CHECK(_request != nullptr);
    _first_block_consumed = false;
    _reader_eof = false;
    _eof = false;
    return Status::OK();
}

Status NativeReader::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 (_request == nullptr) {
        return Status::InternalError("Native v2 reader is not open");
    }

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

    std::string buffer;
    bool local_eof = false;
    if (_first_block_loaded && !_first_block_consumed) {
        buffer = _first_block_buffer;
    } else {
        SCOPED_TIMER(_read_block_time);
        RETURN_IF_ERROR(_read_next_pblock(&buffer, &local_eof));
    }

    if (local_eof && buffer.empty()) {
        _reader_eof = true;
        *eof = true;
        _eof = true;
        return Status::OK();
    }
    if (buffer.empty()) {
        return Status::InternalError("read empty native block from file {}",
                                     _file_description->path);
    }

    PBlock pblock;
    RETURN_IF_ERROR(parse_native_pblock(buffer, _file_description->path, &pblock));
    if (!_schema_inited) {
        RETURN_IF_ERROR(_init_schema_from_pblock(pblock));
    }

    Block source_block;
    size_t uncompressed_bytes = 0;
    int64_t decompress_time = 0;
    {
        SCOPED_TIMER(_deserialize_time);
        RETURN_IF_ERROR(source_block.deserialize(pblock, &uncompressed_bytes, &decompress_time));
    }
    {
        SCOPED_TIMER(_materialize_time);
        RETURN_IF_ERROR(_materialize_requested_columns(source_block, file_block));
    }
    *rows = file_block->rows();
    _record_scan_rows(cast_set<int64_t>(*rows));
    {
        SCOPED_TIMER(_filter_time);
        RETURN_IF_ERROR(_apply_filters(file_block, rows));
    }

    if (_first_block_loaded && !_first_block_consumed) {
        _first_block_consumed = true;
    }
    if (_current_offset >= _file_size) {
        _reader_eof = true;
    }
    *eof = _reader_eof && *rows == 0;
    _eof = *eof;
    return Status::OK();
}

Status NativeReader::close() {
    SCOPED_TIMER(_total_time);
    _file_reader.reset();
    _tracing_file_reader.reset();
    _request.reset();
    _reader_eof = true;
    _eof = true;
    return Status::OK();
}

Status NativeReader::_validate_and_consume_header() {
    DORIS_CHECK(_tracing_file_reader != nullptr);
    _file_size = _tracing_file_reader->size();
    _current_offset = 0;
    _reader_eof = (_file_size == 0);

    static constexpr size_t HEADER_SIZE = sizeof(DORIS_NATIVE_MAGIC) + sizeof(uint32_t);
    if (_reader_eof || _file_size < cast_set<int64_t>(HEADER_SIZE)) {
        return Status::InternalError(
                "invalid Doris Native file {}, file size {} is smaller than header size {}",
                _file_description->path, _file_size, HEADER_SIZE);
    }

    char header[HEADER_SIZE];
    Slice header_slice(header, sizeof(header));
    size_t bytes_read = 0;
    RETURN_IF_ERROR(_tracing_file_reader->read_at(0, header_slice, &bytes_read, _io_ctx.get()));
    if (bytes_read != sizeof(header)) {
        return Status::InternalError(
                "failed to read Doris Native header from file {}, expect {} bytes, got {} bytes",
                _file_description->path, sizeof(header), bytes_read);
    }
    if (std::memcmp(header, DORIS_NATIVE_MAGIC, sizeof(DORIS_NATIVE_MAGIC)) != 0) {
        return Status::InternalError("invalid Doris Native magic header in file {}",
                                     _file_description->path);
    }

    uint32_t version = 0;
    std::memcpy(&version, header + sizeof(DORIS_NATIVE_MAGIC), sizeof(uint32_t));
    if (version != DORIS_NATIVE_FORMAT_VERSION) {
        return Status::InternalError(
                "unsupported Doris Native format version {} in file {}, expect {}", version,
                _file_description->path, DORIS_NATIVE_FORMAT_VERSION);
    }

    _current_offset = sizeof(header);
    _reader_eof = (_file_size == _current_offset);
    return Status::OK();
}

Status NativeReader::_ensure_schema_loaded() const {
    if (_schema_inited) {
        return Status::OK();
    }
    if (!_first_block_loaded) {
        bool local_eof = false;
        RETURN_IF_ERROR(_read_next_pblock(&_first_block_buffer, &local_eof));
        if (local_eof && _first_block_buffer.empty()) {
            return Status::EndOfFile("empty native file {}", _file_description->path);
        }
        if (_first_block_buffer.empty()) {
            return Status::InternalError("first native block is empty {}", _file_description->path);
        }
        _first_block_loaded = true;
    }

    PBlock pblock;
    RETURN_IF_ERROR(parse_native_pblock(_first_block_buffer, _file_description->path, &pblock));
    RETURN_IF_ERROR(_init_schema_from_pblock(pblock));
    return Status::OK();
}

Status NativeReader::_read_next_pblock(std::string* buffer, bool* eof) const {
    DORIS_CHECK(buffer != nullptr);
    DORIS_CHECK(eof != nullptr);
    DORIS_CHECK(_tracing_file_reader != nullptr);
    buffer->clear();
    *eof = false;

    if (_current_offset >= _file_size) {
        *eof = true;
        return Status::OK();
    }

    const auto remaining_prefix_bytes = _file_size - _current_offset;
    if (remaining_prefix_bytes < sizeof(uint64_t)) {
        // Header-only is the one valid empty-file representation and returned above. Once any
        // prefix byte exists, all eight bytes are mandatory. For example, `header + 4 bytes` is a
        // truncated Native file, not EOF that FileScannerV2 may skip as an empty split.
        return Status::InternalError(
                "truncated native block length in file {} at offset {}, expect {}, remaining {}",
                _file_description->path, _current_offset, sizeof(uint64_t), remaining_prefix_bytes);
    }

    uint64_t block_len = 0;
    Slice len_slice(reinterpret_cast<char*>(&block_len), sizeof(block_len));
    size_t bytes_read = 0;
    RETURN_IF_ERROR(
            _tracing_file_reader->read_at(_current_offset, len_slice, &bytes_read, _io_ctx.get()));
    if (bytes_read == 0) {
        *eof = true;
        return Status::OK();
    }
    if (bytes_read != sizeof(block_len)) {
        return Status::InternalError(
                "Failed to read native block length from file {}, expect {}, actual {}",
                _file_description->path, sizeof(block_len), bytes_read);
    }
    _current_offset += sizeof(block_len);
    if (block_len == 0) {
        // A header-only file reaches the `_current_offset >= _file_size` branch above and is a
        // valid empty Native file. Once an explicit length prefix is present, however, zero is not
        // an EOF marker in the Native format. For example, `header + uint64(0)` is truncated input,
        // not a zero-row split, and must not escape as EOF for FileScannerV2 to count as empty.
        return Status::InternalError("zero-length native block in file {} at offset {}",
                                     _file_description->path, _current_offset - sizeof(block_len));
    }

    const auto remaining_block_bytes = cast_set<uint64_t>(_file_size - _current_offset);
    if (block_len > remaining_block_bytes) {
        // Validate before allocating `block_len` bytes. Besides reporting a precise corruption
        // error, this prevents a truncated file with a damaged length prefix from requesting an
        // allocation much larger than the physical file.
        return Status::InternalError(
                "truncated native block body in file {} at offset {}, expect {}, remaining {}",
                _file_description->path, _current_offset, block_len, remaining_block_bytes);
    }

    buffer->assign(block_len, '\0');
    Slice data_slice(buffer->data(), block_len);
    bytes_read = 0;
    RETURN_IF_ERROR(
            _tracing_file_reader->read_at(_current_offset, data_slice, &bytes_read, _io_ctx.get()));
    if (bytes_read != block_len) {
        return Status::InternalError(
                "Failed to read native block body from file {}, expect {}, actual {}",
                _file_description->path, block_len, bytes_read);
    }
    _current_offset += block_len;
    *eof = (_current_offset >= _file_size);
    return Status::OK();
}

Status NativeReader::_init_schema_from_pblock(const PBlock& pblock) const {
    _file_schema.clear();
    _file_schema.reserve(pblock.column_metas_size());
    for (int idx = 0; idx < pblock.column_metas_size(); ++idx) {
        const auto& meta = pblock.column_metas(idx);
        ColumnDefinition field;
        field.identifier = Field::create_field<TYPE_STRING>(meta.name());
        field.local_id = idx;
        field.name = meta.name();
        field.type = make_nullable(DataTypeFactory::instance().create_data_type(meta));
        _file_schema.push_back(std::move(field));
    }
    _schema_inited = true;
    return Status::OK();
}

Status NativeReader::_materialize_requested_columns(const Block& source_block,
                                                    Block* file_block) const {
    DORIS_CHECK(file_block != nullptr);
    DORIS_CHECK(_request != 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::InternalError("native file {} does not contain local column id {}",
                                         _file_description->path, source_idx);
        }
        if (block_position.value() >= file_block->columns()) {
            return Status::InternalError("native v2 request has invalid block position {}",
                                         block_position.value());
        }
        const auto& target = file_block->get_by_position(block_position.value());
        auto column = 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();
}

Status NativeReader::_apply_filters(Block* file_block, size_t* rows) const {
    return apply_materialized_reader_filters(_request.get(), _io_ctx.get(), file_block, rows);
}

} // namespace doris::format::native