// 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 "writer.h"

#include <type_traits>

#include "pipeline/exec/exchange_sink_operator.h"

#include "vec/core/block.h"

namespace doris::pipeline {

#include "common/compile_check_begin.h"

template <typename ChannelPtrType>

void Writer::_handle_eof_channel(RuntimeState* state, ChannelPtrType channel, Status st) const {

    channel->set_receiver_eof(st);

    // Chanel will not send RPC to the downstream when eof, so close chanel by OK status.

    static_cast<void>(channel->close(state));

}

Status Writer::write(ExchangeSinkLocalState* local_state, RuntimeState* state,

                     vectorized::Block* block, bool eos) {

    bool already_sent = false;

    {

        SCOPED_TIMER(local_state->split_block_hash_compute_timer());

        RETURN_IF_ERROR(

                local_state->partitioner()->do_partitioning(state, block, eos, &already_sent));

    }

    if (already_sent) {

        // The same block may be sent twice by TabletSinkHashPartitioner. To get the correct

        // result, we should not send any rows the last time.

        return Status::OK();

    }

    auto rows = block->rows();

    {

        SCOPED_TIMER(local_state->distribute_rows_into_channels_timer());

        const auto& channel_filed = local_state->partitioner()->get_channel_ids();

        if (channel_filed.len == sizeof(uint32_t)) {

            RETURN_IF_ERROR(_channel_add_rows(state, local_state->channels,

                                              local_state->channels.size(),

                                              channel_filed.get<uint32_t>(), rows, block, eos));

        } else {

            RETURN_IF_ERROR(_channel_add_rows(state, local_state->channels,

                                              local_state->channels.size(),

                                              channel_filed.get<int64_t>(), rows, block, eos));

        }

    }

    return Status::OK();

}

template <typename ChannelIdType>

Status Writer::_channel_add_rows(RuntimeState* state,

                                 std::vector<std::shared_ptr<vectorized::Channel>>& channels,

                                 size_t partition_count,

                                 const ChannelIdType* __restrict channel_ids, size_t rows,

                                 vectorized::Block* block, bool eos) {

    _row_idx.resize(rows);

    {

        _partition_rows_histogram.resize(partition_count);

        _channel_start_offsets.resize(partition_count);

        for (size_t i = 0; i < partition_count; ++i) {

            _partition_rows_histogram[i] = 0;

        }

        for (size_t i = 0; i < rows; ++i) {

            _partition_rows_histogram[channel_ids[i]]++;

        }

        _channel_start_offsets[0] = 0;

        for (size_t i = 1; i < partition_count; ++i) {

            _channel_start_offsets[i] =

                    _channel_start_offsets[i - 1] + _partition_rows_histogram[i - 1];

        }

        for (uint32_t i = 0; i < rows; i++) {

            if constexpr (std::is_signed_v<ChannelIdType>) {

                // -1 means this row is filtered by table sink hash partitioner

                if (channel_ids[i] == -1) {

                    continue;

                }

            }

            _row_idx[_channel_start_offsets[channel_ids[i]]++] = i;

        }

    }

    Status status = Status::OK();

    uint32_t offset = 0;

    for (size_t i = 0; i < partition_count; ++i) {

        uint32_t size = _partition_rows_histogram[i];

        if (!channels[i]->is_receiver_eof() && size > 0) {

            status = channels[i]->add_rows(block, _row_idx.data(), offset, size, false);

            HANDLE_CHANNEL_STATUS(state, channels[i], status);

        }

        offset += size;

    }

    if (eos) {

        for (int i = 0; i < partition_count; ++i) {

            if (!channels[i]->is_receiver_eof()) {

                status = channels[i]->add_rows(block, _row_idx.data(), 0, 0, true);

                HANDLE_CHANNEL_STATUS(state, channels[i], status);

            }

        }

    }

    return Status::OK();

}

Unexecuted instantiation: _ZN5doris8pipeline6Writer17_channel_add_rowsIjEENS_6StatusEPNS_12RuntimeStateERSt6vectorISt10shared_ptrINS_10vectorized7ChannelEESaISA_EEmPKT_mPNS8_5BlockEb

Unexecuted instantiation: _ZN5doris8pipeline6Writer17_channel_add_rowsIlEENS_6StatusEPNS_12RuntimeStateERSt6vectorISt10shared_ptrINS_10vectorized7ChannelEESaISA_EEmPKT_mPNS8_5BlockEb

} // namespace doris::pipeline