be/src/exec/exchange/local_exchanger.cpp

Source
// 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 "exec/exchange/local_exchanger.h"

#include "common/cast_set.h"
#include "common/status.h"
#include "exec/exchange/local_exchange_sink_operator.h"
#include "exec/exchange/local_exchange_source_operator.h"
#include "exec/partitioner/partitioner.h"

namespace doris {
template <typename BlockType>
void Exchanger<BlockType>::_enqueue_data_and_set_ready(int channel_id,
                                                       LocalExchangeSinkLocalState* local_state,
                                                       BlockType&& block) {
    if (local_state == nullptr) {
        _enqueue_data_and_set_ready(channel_id, std::move(block));
        return;
    }
    // PartitionedBlock is used by shuffle exchanger.
    // PartitionedBlock will be push into multiple queues with different row ranges, so it will be
    // referenced multiple times. Otherwise, we only ref the block once because it is only push into
    // one queue.
    std::unique_lock l(*_m[channel_id]);
    if constexpr (std::is_same_v<PartitionedBlock, BlockType> ||
                  std::is_same_v<BroadcastBlock, BlockType>) {
        block.first->record_channel_id(channel_id);
    } else {
        block->record_channel_id(channel_id);
    }

    if (_data_queue[channel_id].enqueue(std::move(block))) {
        local_state->_shared_state->set_ready_to_read(channel_id);
    }
}

template <typename BlockType>
bool Exchanger<BlockType>::_dequeue_data(LocalExchangeSourceLocalState* local_state,
                                         BlockType& block, bool* eos, Block* data_block,
                                         int channel_id) {
    if (local_state == nullptr) {
        return _dequeue_data(block, eos, data_block, channel_id);
    }
    bool all_finished = _running_sink_operators == 0;
    if (_data_queue[channel_id].try_dequeue(block)) {
        if constexpr (std::is_same_v<PartitionedBlock, BlockType> ||
                      std::is_same_v<BroadcastBlock, BlockType>) {
            local_state->_shared_state->sub_mem_usage(channel_id, block.first->_allocated_bytes);
        } else {
            local_state->_shared_state->sub_mem_usage(channel_id, block->_allocated_bytes);
            data_block->swap(block->_data_block);
        }
        return true;
    } else if (all_finished) {
        *eos = true;
    } else {
        std::unique_lock l(*_m[channel_id]);
        if (_data_queue[channel_id].try_dequeue(block)) {
            if constexpr (std::is_same_v<PartitionedBlock, BlockType> ||
                          std::is_same_v<BroadcastBlock, BlockType>) {
                local_state->_shared_state->sub_mem_usage(channel_id,
                                                          block.first->_allocated_bytes);
            } else {
                local_state->_shared_state->sub_mem_usage(channel_id, block->_allocated_bytes);
                data_block->swap(block->_data_block);
            }
            return true;
        }
        COUNTER_UPDATE(local_state->_get_block_failed_counter, 1);
        local_state->_dependency->block();
    }
    return false;
}

template <typename BlockType>
void Exchanger<BlockType>::_enqueue_data_and_set_ready(int channel_id, BlockType&& block) {
    if constexpr (std::is_same_v<PartitionedBlock, BlockType> ||
                  std::is_same_v<BroadcastBlock, BlockType>) {
        block.first->record_channel_id(channel_id);
    } else {
        block->record_channel_id(channel_id);
    }
    _data_queue[channel_id].enqueue(std::move(block));
}

template <typename BlockType>
bool Exchanger<BlockType>::_dequeue_data(BlockType& block, bool* eos, Block* data_block,
                                         int channel_id) {
    if (_data_queue[channel_id].try_dequeue(block)) {
        if constexpr (!std::is_same_v<PartitionedBlock, BlockType> &&
                      !std::is_same_v<BroadcastBlock, BlockType>) {
            data_block->swap(block->_data_block);
        }
        return true;
    }
    return false;
}

Status ShuffleExchanger::sink(RuntimeState* state, Block* in_block, bool eos, Profile&& profile,
                              SinkInfo& sink_info) {
    if (in_block->empty()) {
        return Status::OK();
    }
    {
        SCOPED_TIMER(profile.compute_hash_value_timer);
        RETURN_IF_ERROR(sink_info.partitioner->do_partitioning(state, in_block));
    }
    {
        SCOPED_TIMER(profile.distribute_timer);
        RETURN_IF_ERROR(_split_rows(state, sink_info.partitioner->get_channel_ids(), in_block,
                                    *sink_info.channel_id, sink_info.local_state,
                                    sink_info.shuffle_idx_to_instance_idx));
    }

    sink_info.local_state->_memory_used_counter->set(
            sink_info.local_state->_shared_state->mem_usage);
    return Status::OK();
}

void ShuffleExchanger::close(SourceInfo&& source_info) {
    PartitionedBlock partitioned_block;
    bool eos;
    Block block;
    _data_queue[source_info.channel_id].set_eos();
    while (_dequeue_data(source_info.local_state, partitioned_block, &eos, &block,
                         source_info.channel_id)) {
        // do nothing
    }
}

Status ShuffleExchanger::get_block(RuntimeState* state, Block* block, bool* eos, Profile&& profile,
                                   SourceInfo&& source_info) {
    PartitionedBlock partitioned_block;
    MutableBlock mutable_block;

    auto get_data = [&]() -> Status {
        do {
            const auto* offset_start = partitioned_block.second.row_idxs->data() +
                                       partitioned_block.second.offset_start;
            auto block_wrapper = partitioned_block.first;
            RETURN_IF_ERROR(mutable_block.add_rows(&block_wrapper->_data_block, offset_start,
                                                   offset_start + partitioned_block.second.length));
        } while (mutable_block.rows() < state->batch_size() && !*eos &&
                 _dequeue_data(source_info.local_state, partitioned_block, eos, block,
                               source_info.channel_id));
        return Status::OK();
    };

    if (_dequeue_data(source_info.local_state, partitioned_block, eos, block,
                      source_info.channel_id)) {
        SCOPED_TIMER(profile.copy_data_timer);
        mutable_block = VectorizedUtils::build_mutable_mem_reuse_block(
                block, partitioned_block.first->_data_block);
        RETURN_IF_ERROR(get_data());
    }
    return Status::OK();
}

Status ShuffleExchanger::_split_rows(RuntimeState* state, const std::vector<uint32_t>& channel_ids,
                                     Block* block, int channel_id,
                                     LocalExchangeSinkLocalState* local_state,
                                     std::map<int, int>* shuffle_idx_to_instance_idx) {
    if (local_state == nullptr) {
        return _split_rows(state, channel_ids, block, channel_id);
    }
    const auto rows = cast_set<int32_t>(block->rows());
    auto row_idx = std::make_shared<PODArray<uint32_t>>(rows);
    auto& partition_rows_histogram = _partition_rows_histogram[channel_id];
    {
        partition_rows_histogram.assign(_num_partitions + 1, 0);
        for (int32_t i = 0; i < rows; ++i) {
            partition_rows_histogram[channel_ids[i]]++;
        }
        for (int32_t i = 1; i <= _num_partitions; ++i) {
            partition_rows_histogram[i] += partition_rows_histogram[i - 1];
        }
        for (int32_t i = rows - 1; i >= 0; --i) {
            (*row_idx)[partition_rows_histogram[channel_ids[i]] - 1] = i;
            partition_rows_histogram[channel_ids[i]]--;
        }
    }

    Block data_block;
    std::shared_ptr<BlockWrapper> new_block_wrapper;
    if (!_free_blocks.try_dequeue(data_block)) {
        data_block = block->clone_empty();
    }
    data_block.swap(*block);
    new_block_wrapper =
            BlockWrapper::create_shared(std::move(data_block), local_state->_shared_state, -1);
    if (new_block_wrapper->_data_block.empty()) {
        return Status::OK();
    }
    /**
     * Data are hash-shuffled and distributed to all instances of
     * all BEs. So we need a shuffleId-To-InstanceId mapping.
     * For example, row 1 get a hash value 1 which means we should distribute to instance 1 on
     * BE 1 and row 2 get a hash value 2 which means we should distribute to instance 1 on BE 3.
     */
    DCHECK(shuffle_idx_to_instance_idx && shuffle_idx_to_instance_idx->size() > 0);
    const auto& map = *shuffle_idx_to_instance_idx;
    int32_t enqueue_rows = 0;
    for (const auto& it : map) {
        DCHECK(it.second >= 0 && it.second < _num_partitions)
                << it.first << " : " << it.second << " " << _num_partitions;
        uint32_t start = partition_rows_histogram[it.first];
        uint32_t size = partition_rows_histogram[it.first + 1] - start;
        if (size > 0) {
            enqueue_rows += size;
            _enqueue_data_and_set_ready(
                    it.second, local_state,
                    {new_block_wrapper,
                     {.row_idxs = row_idx, .offset_start = start, .length = size}});
        }
    }
    if (enqueue_rows != rows) [[unlikely]] {
        fmt::memory_buffer debug_string_buffer;
        fmt::format_to(debug_string_buffer, "Type: {}, Local Exchange Id: {}, Shuffled Map: ",
                       get_exchange_type_name(get_type()), local_state->parent()->node_id());
        for (const auto& it : map) {
            fmt::format_to(debug_string_buffer, "[{}:{}], ", it.first, it.second);
        }
        return Status::InternalError(
                "Rows mismatched! Data may be lost. [Expected enqueue rows={}, Real enqueue "
                "rows={}, Detail: {}]",
                rows, enqueue_rows, fmt::to_string(debug_string_buffer));
    }

    return Status::OK();
}

Status ShuffleExchanger::_split_rows(RuntimeState* state, const std::vector<uint32_t>& channel_ids,
                                     Block* block, int channel_id) {
    const auto rows = cast_set<int32_t>(block->rows());
    auto row_idx = std::make_shared<PODArray<uint32_t>>(rows);
    auto& partition_rows_histogram = _partition_rows_histogram[channel_id];
    {
        partition_rows_histogram.assign(_num_partitions + 1, 0);
        for (int32_t i = 0; i < rows; ++i) {
            partition_rows_histogram[channel_ids[i]]++;
        }
        for (int32_t i = 1; i <= _num_partitions; ++i) {
            partition_rows_histogram[i] += partition_rows_histogram[i - 1];
        }
        for (int32_t i = rows - 1; i >= 0; --i) {
            (*row_idx)[partition_rows_histogram[channel_ids[i]] - 1] = i;
            partition_rows_histogram[channel_ids[i]]--;
        }
    }

    Block data_block;
    std::shared_ptr<BlockWrapper> new_block_wrapper;
    if (!_free_blocks.try_dequeue(data_block)) {
        data_block = block->clone_empty();
    }
    data_block.swap(*block);
    new_block_wrapper = BlockWrapper::create_shared(std::move(data_block), nullptr, -1);
    if (new_block_wrapper->_data_block.empty()) {
        return Status::OK();
    }
    for (int i = 0; i < _num_partitions; i++) {
        uint32_t start = partition_rows_histogram[i];
        uint32_t size = partition_rows_histogram[i + 1] - start;
        if (size > 0) {
            _enqueue_data_and_set_ready(
                    i, {new_block_wrapper,
                        {.row_idxs = row_idx, .offset_start = start, .length = size}});
        }
    }

    return Status::OK();
}

Status PassthroughExchanger::sink(RuntimeState* state, Block* in_block, bool eos, Profile&& profile,
                                  SinkInfo& sink_info) {
    if (in_block->empty()) {
        return Status::OK();
    }
    Block new_block;
    if (!_free_blocks.try_dequeue(new_block)) {
        new_block = {in_block->clone_empty()};
    }
    new_block.swap(*in_block);
    auto channel_id = ((*sink_info.channel_id)++) % _num_partitions;
    BlockWrapperSPtr wrapper = BlockWrapper::create_shared(
            std::move(new_block),
            sink_info.local_state ? sink_info.local_state->_shared_state : nullptr, channel_id);

    _enqueue_data_and_set_ready(channel_id, sink_info.local_state, std::move(wrapper));

    sink_info.local_state->_memory_used_counter->set(
            sink_info.local_state->_shared_state->mem_usage);

    return Status::OK();
}

void PassthroughExchanger::close(SourceInfo&& source_info) {
    Block next_block;
    BlockWrapperSPtr wrapper;
    bool eos;
    _data_queue[source_info.channel_id].set_eos();
    while (_dequeue_data(source_info.local_state, wrapper, &eos, &next_block,
                         source_info.channel_id)) {
        // do nothing
    }
}

void PassToOneExchanger::close(SourceInfo&& source_info) {
    Block next_block;
    BlockWrapperSPtr wrapper;
    bool eos;
    _data_queue[source_info.channel_id].set_eos();
    while (_dequeue_data(source_info.local_state, wrapper, &eos, &next_block,
                         source_info.channel_id)) {
        // do nothing
    }
}

Status PassthroughExchanger::get_block(RuntimeState* state, Block* block, bool* eos,
                                       Profile&& profile, SourceInfo&& source_info) {
    BlockWrapperSPtr next_block;
    _dequeue_data(source_info.local_state, next_block, eos, block, source_info.channel_id);
    return Status::OK();
}

Status PassToOneExchanger::sink(RuntimeState* state, Block* in_block, bool eos, Profile&& profile,
                                SinkInfo& sink_info) {
    if (in_block->empty()) {
        return Status::OK();
    }
    Block new_block;
    if (!_free_blocks.try_dequeue(new_block)) {
        new_block = {in_block->clone_empty()};
    }
    new_block.swap(*in_block);

    BlockWrapperSPtr wrapper = BlockWrapper::create_shared(
            std::move(new_block),
            sink_info.local_state ? sink_info.local_state->_shared_state : nullptr, 0);
    _enqueue_data_and_set_ready(0, sink_info.local_state, std::move(wrapper));

    sink_info.local_state->_memory_used_counter->set(
            sink_info.local_state->_shared_state->mem_usage);

    return Status::OK();
}

Status PassToOneExchanger::get_block(RuntimeState* state, Block* block, bool* eos,
                                     Profile&& profile, SourceInfo&& source_info) {
    if (source_info.channel_id != 0) {
        *eos = true;
        return Status::OK();
    }
    BlockWrapperSPtr next_block;
    _dequeue_data(source_info.local_state, next_block, eos, block, source_info.channel_id);
    return Status::OK();
}

void ExchangerBase::finalize() {
    DCHECK(_running_source_operators == 0);
    Block block;
    while (_free_blocks.try_dequeue(block)) {
        // do nothing
    }
}

Status BroadcastExchanger::sink(RuntimeState* state, Block* in_block, bool eos, Profile&& profile,
                                SinkInfo& sink_info) {
    if (in_block->empty()) {
        return Status::OK();
    }
    Block new_block;
    if (!_free_blocks.try_dequeue(new_block)) {
        new_block = {in_block->clone_empty()};
    }
    new_block.swap(*in_block);
    auto wrapper = BlockWrapper::create_shared(
            std::move(new_block),
            sink_info.local_state ? sink_info.local_state->_shared_state : nullptr, -1);
    for (int i = 0; i < _num_partitions; i++) {
        _enqueue_data_and_set_ready(
                i, sink_info.local_state,
                {wrapper, {.offset_start = 0, .length = wrapper->_data_block.rows()}});
    }

    return Status::OK();
}

void BroadcastExchanger::close(SourceInfo&& source_info) {
    BroadcastBlock partitioned_block;
    bool eos;
    Block block;
    _data_queue[source_info.channel_id].set_eos();
    while (_dequeue_data(source_info.local_state, partitioned_block, &eos, &block,
                         source_info.channel_id)) {
        // do nothing
    }
}

Status BroadcastExchanger::get_block(RuntimeState* state, Block* block, bool* eos,
                                     Profile&& profile, SourceInfo&& source_info) {
    BroadcastBlock partitioned_block;

    if (_dequeue_data(source_info.local_state, partitioned_block, eos, block,
                      source_info.channel_id)) {
        SCOPED_TIMER(profile.copy_data_timer);
        MutableBlock mutable_block = VectorizedUtils::build_mutable_mem_reuse_block(
                block, partitioned_block.first->_data_block);
        auto block_wrapper = partitioned_block.first;
        RETURN_IF_ERROR(mutable_block.add_rows(&block_wrapper->_data_block,
                                               partitioned_block.second.offset_start,
                                               partitioned_block.second.length));
    }

    return Status::OK();
}

Status AdaptivePassthroughExchanger::_passthrough_sink(RuntimeState* state, Block* in_block,
                                                       SinkInfo& sink_info) {
    Block new_block;
    if (!_free_blocks.try_dequeue(new_block)) {
        new_block = {in_block->clone_empty()};
    }
    new_block.swap(*in_block);
    auto channel_id = ((*sink_info.channel_id)++) % _num_partitions;
    _enqueue_data_and_set_ready(
            channel_id, sink_info.local_state,
            {BlockWrapper::create_shared(
                     std::move(new_block),
                     sink_info.local_state ? sink_info.local_state->_shared_state : nullptr, -1),
             {.row_idxs = nullptr, .offset_start = 0, .length = 0}});

    sink_info.local_state->_memory_used_counter->set(
            sink_info.local_state->_shared_state->mem_usage);
    return Status::OK();
}

Status AdaptivePassthroughExchanger::_shuffle_sink(RuntimeState* state, Block* block,
                                                   SinkInfo& sink_info) {
    std::vector<uint32_t> channel_ids;
    const auto num_rows = block->rows();
    channel_ids.resize(num_rows, 0);
    if (num_rows <= _num_partitions) {
        std::iota(channel_ids.begin(), channel_ids.end(), 0);
    } else {
        size_t i = 0;
        for (; i < num_rows - _num_partitions; i += _num_partitions) {
            std::iota(channel_ids.begin() + i, channel_ids.begin() + i + _num_partitions, 0);
        }
        if (i < num_rows - 1) {
            std::iota(channel_ids.begin() + i, channel_ids.end(), 0);
        }
    }

    sink_info.local_state->_memory_used_counter->set(
            sink_info.local_state->_shared_state->mem_usage);
    RETURN_IF_ERROR(_split_rows(state, channel_ids, block, sink_info));
    return Status::OK();
}

Status AdaptivePassthroughExchanger::_split_rows(RuntimeState* state,
                                                 const std::vector<uint32_t>& channel_ids,
                                                 Block* block, SinkInfo& sink_info) {
    const auto rows = cast_set<int32_t>(block->rows());
    auto row_idx = std::make_shared<PODArray<uint32_t>>(rows);
    auto& partition_rows_histogram = _partition_rows_histogram[sink_info.ins_idx];
    {
        partition_rows_histogram.assign(_num_partitions + 1, 0);
        for (int32_t i = 0; i < rows; ++i) {
            partition_rows_histogram[channel_ids[i]]++;
        }
        for (int32_t i = 1; i <= _num_partitions; ++i) {
            partition_rows_histogram[i] += partition_rows_histogram[i - 1];
        }

        for (int32_t i = rows - 1; i >= 0; --i) {
            (*row_idx)[partition_rows_histogram[channel_ids[i]] - 1] = i;
            partition_rows_histogram[channel_ids[i]]--;
        }
    }
    Block data_block;
    if (!_free_blocks.try_dequeue(data_block)) {
        data_block = block->clone_empty();
    }
    data_block.swap(*block);
    std::shared_ptr<BlockWrapper> new_block_wrapper = BlockWrapper::create_shared(
            std::move(data_block), sink_info.local_state->_shared_state, sink_info.ins_idx);
    if (new_block_wrapper->_data_block.empty()) {
        return Status::OK();
    }
    for (int32_t i = 0; i < _num_partitions; i++) {
        const uint32_t start = partition_rows_histogram[i];
        const uint32_t size = partition_rows_histogram[i + 1] - start;
        if (size > 0) {
            _enqueue_data_and_set_ready(
                    i, sink_info.local_state,
                    {new_block_wrapper,
                     {.row_idxs = row_idx, .offset_start = start, .length = size}});
        }
    }
    return Status::OK();
}

Status AdaptivePassthroughExchanger::sink(RuntimeState* state, Block* in_block, bool eos,
                                          Profile&& profile, SinkInfo& sink_info) {
    if (in_block->empty()) {
        return Status::OK();
    }
    if (_is_pass_through) {
        return _passthrough_sink(state, in_block, sink_info);
    } else {
        if (++_total_block >= _num_partitions) {
            _is_pass_through = true;
        }
        return _shuffle_sink(state, in_block, sink_info);
    }
}

Status AdaptivePassthroughExchanger::get_block(RuntimeState* state, Block* block, bool* eos,
                                               Profile&& profile, SourceInfo&& source_info) {
    if (!_tmp_block[source_info.channel_id].empty()) {
        *block = std::move(_tmp_block[source_info.channel_id]);
        *eos = _tmp_eos[source_info.channel_id];
        _tmp_block[source_info.channel_id] = {};
        return Status::OK();
    }
    PartitionedBlock partitioned_block;
    MutableBlock mutable_block;

    auto get_data = [&]() -> Status {
        do {
            if (partitioned_block.second.row_idxs == nullptr) {
                // The passthrough path which means the block is not partitioned, we can directly move the block without copying.
                if (mutable_block.rows() > 0) {
                    _tmp_block[source_info.channel_id] =
                            std::move(partitioned_block.first->_data_block);
                    _tmp_eos[source_info.channel_id] = *eos;
                    *eos = false;
                } else {
                    *block = std::move(partitioned_block.first->_data_block);
                }
                break;
            }
            const auto* offset_start = partitioned_block.second.row_idxs->data() +
                                       partitioned_block.second.offset_start;
            auto block_wrapper = partitioned_block.first;
            RETURN_IF_ERROR(mutable_block.add_rows(&block_wrapper->_data_block, offset_start,
                                                   offset_start + partitioned_block.second.length));
        } while (mutable_block.rows() < state->batch_size() && !*eos &&
                 _dequeue_data(source_info.local_state, partitioned_block, eos, block,
                               source_info.channel_id));
        return Status::OK();
    };

    if (_dequeue_data(source_info.local_state, partitioned_block, eos, block,
                      source_info.channel_id)) {
        SCOPED_TIMER(profile.copy_data_timer);
        mutable_block = VectorizedUtils::build_mutable_mem_reuse_block(
                block, partitioned_block.first->_data_block);
        RETURN_IF_ERROR(get_data());
    }
    return Status::OK();
}

void AdaptivePassthroughExchanger::close(SourceInfo&& source_info) {
    PartitionedBlock partitioned_block;
    bool eos;
    Block block;
    _data_queue[source_info.channel_id].set_eos();
    while (_dequeue_data(source_info.local_state, partitioned_block, &eos, &block,
                         source_info.channel_id)) {
        // do nothing
    }
}

} // namespace doris

Coverage Report

Created: 2026-05-13 20:09