Function GetRandomDatasetBatch

catboost/libs/fstr/sage_values.cpp:230–277

}

TDataProvider GetRandomDatasetBatch(
    const TDataProvider& dataset,
    size_t batchSize,
    TRestorableFastRng64* randPtr,
    NPar::ILocalExecutor* localExecutor)
{
    if (dataset.ObjectsGrouping->IsTrivial()) {
        TVector<ui32> indices(dataset.GetObjectCount());
        Iota(indices.begin(), indices.end(), 0);
        PartialShuffle(indices.begin(), indices.end(), batchSize, *randPtr);

        auto subset = dataset.GetSubset(
            GetSubset(
                dataset.ObjectsGrouping.Get(),
                std::move(TArraySubsetIndexing<ui32>(TVector<ui32>(indices.begin(), indices.begin() + batchSize))),
                EObjectsOrder::Ordered
            ),
            GetMonopolisticFreeCpuRam(),
            localExecutor
        );

        return *subset;
    } else {
        TVector<ui32> groupIndices(dataset.ObjectsGrouping->GetGroupCount());
        Iota(groupIndices.begin(), groupIndices.end(), 0);
        Shuffle(groupIndices.begin(), groupIndices.end(), *randPtr);

        TVector<ui32> batchGroupIndices;
        ui32 realBatchSize = 0;
        ui32 groupShuffledIndex = 0;
        while (realBatchSize < batchSize) {
            realBatchSize += dataset.ObjectsGrouping->GetGroup(groupIndices[groupShuffledIndex]).GetSize();
            ++groupShuffledIndex;
        }

        auto subset = dataset.GetSubset(
            GetSubset(
                dataset.ObjectsGrouping.Get(),
                std::move(TArraySubsetIndexing<ui32>(TVector<ui32>(groupIndices.begin(), groupIndices.begin() + groupShuffledIndex))),
                EObjectsOrder::Ordered
            ),
            GetMonopolisticFreeCpuRam(),
            localExecutor
        );

        return *subset;
    }
}

double CalculateModelLoss(
    const TFullModel& model,