Method GetSubset

catboost/libs/data/objects_grouping.cpp:27–123

TObjectsGroupingSubset NCB::GetSubset(
    TObjectsGroupingPtr objectsGrouping,
    TArraySubsetIndexing<ui32>&& groupsSubset,
    EObjectsOrder groupSubsetOrder
) {
    if (objectsGrouping->IsTrivial()) {
        return TObjectsGroupingSubset(
            ::std::holds_alternative<TFullSubset<ui32>>(groupsSubset) ?
                objectsGrouping : MakeIntrusive<TObjectsGrouping>(groupsSubset.Size()),
            std::move(groupsSubset),
            groupSubsetOrder
        );
    } else {
        TMaybe<TArraySubsetIndexing<ui32>> objectsSubset;
        TVector<TGroupBounds> subsetGroupBounds;

        if (std::holds_alternative<TFullSubset<ui32>>(groupsSubset)) {
            objectsSubset = MakeMaybe<TArraySubsetIndexing<ui32>>(
                TFullSubset<ui32>(objectsGrouping->GetObjectCount())
            );
            return TObjectsGroupingSubset(
                objectsGrouping,
                std::move(groupsSubset),
                groupSubsetOrder,
                MakeMaybe<TArraySubsetIndexing<ui32>>(
                    TFullSubset<ui32>(objectsGrouping->GetObjectCount())
                ),
                groupSubsetOrder
            );
        } else if (const auto* ranges = std::get_if<TRangesSubset<ui32>>(&groupsSubset)) {
            const auto& groupsSubsetBlocks = ranges->Blocks;

            ui32 srcObjectCount = objectsGrouping->GetObjectCount();
            auto nontrivialSrcGroups = objectsGrouping->GetNonTrivialGroups();

            subsetGroupBounds.reserve(groupsSubset.Size());

            TVector<TSubsetBlock<ui32>> objectsSubsetBlocks;
            objectsSubsetBlocks.reserve(groupsSubsetBlocks.size());

            ui32 objectsDstBegin = 0;
            for (const auto& groupSubsetBlock : groupsSubsetBlocks) {
                objectsSubsetBlocks.emplace_back(
                    TIndexRange<ui32>(
                        nontrivialSrcGroups[groupSubsetBlock.SrcBegin].Begin,
                        (groupSubsetBlock.SrcEnd == nontrivialSrcGroups.size() ?
                            srcObjectCount : nontrivialSrcGroups[groupSubsetBlock.SrcEnd].Begin)
                    ),
                    objectsDstBegin
                );

                for (auto srcGroupIdx : xrange(groupSubsetBlock.SrcBegin, groupSubsetBlock.SrcEnd)) {
                    subsetGroupBounds.emplace_back(
                        objectsDstBegin,
                        objectsDstBegin + nontrivialSrcGroups[srcGroupIdx].GetSize()
                    );
                    objectsDstBegin += subsetGroupBounds.back().GetSize();
                }