be/src/format/count_reader.h

// 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.

#pragma once

#include <cstddef>
#include <cstdint>
#include <memory>

#include "common/status.h"
#include "core/block/block.h"
#include "format/generic_reader.h"

namespace doris {

/// A lightweight reader that emits row counts without reading any actual data.
/// Used as a decorator to replace the real reader when COUNT(*) push down is active.
///
/// Instead of duplicating the COUNT short-circuit logic in every format reader
/// (ORC, Parquet, etc.), FileScanner creates a CountReader after the real reader
/// is initialized and the total row count is known. The CountReader then serves
/// all subsequent get_next_block calls by simply resizing columns.
///
/// This cleanly separates the "how many rows" concern from the actual data reading,
/// eliminating duplicated COUNT blocks across format readers.
class CountReader : public GenericReader {
public:
    /// @param total_rows   Total number of rows to emit (post-filter).
    /// @param batch_size   Maximum rows per batch.
    /// @param inner_reader The original reader, kept alive for profile collection
    ///                     and lifecycle management. Ownership is transferred.
    CountReader(int64_t total_rows, size_t batch_size,
                std::unique_ptr<GenericReader> inner_reader = nullptr)
            : _total_rows(total_rows),
              _remaining_rows(total_rows),
              _batch_size(batch_size),
              _inner_reader(std::move(inner_reader)) {
        set_push_down_agg_type(TPushAggOp::type::COUNT);
    }

    ~CountReader() override = default;

    Status _do_get_next_block(Block* block, size_t* read_rows, bool* eof) override {
        auto rows = std::min(_remaining_rows, static_cast<int64_t>(_batch_size));
        _remaining_rows -= rows;

        auto mutate_columns = block->mutate_columns();
        for (auto& col : mutate_columns) {
            col->resize(rows);
        }
        block->set_columns(std::move(mutate_columns));

        *read_rows = rows;
        *eof = (_remaining_rows == 0);
        return Status::OK();
    }

    /// CountReader counts rows by definition.
    bool count_read_rows() override { return true; }

    /// Delegate to inner reader if available, otherwise return our total.
    int64_t get_total_rows() const override {
        return _inner_reader ? _inner_reader->get_total_rows() : _initial_total_rows();
    }

    Status close() override {
        if (_inner_reader) {
            return _inner_reader->close();
        }
        return Status::OK();
    }

    /// Access the inner reader for profile collection or other lifecycle needs.
    GenericReader* inner_reader() const { return _inner_reader.get(); }

protected:
    void _collect_profile_before_close() override {
        if (_inner_reader) {
            _inner_reader->collect_profile_before_close();
        }
    }

private:
    int64_t _initial_total_rows() const { return _total_rows; }

    const int64_t _total_rows;
    int64_t _remaining_rows;
    size_t _batch_size;
    std::unique_ptr<GenericReader> _inner_reader;
};

} // namespace doris