/*****************************************************************************
 *  $Id: seg_plan.cpp 646372 2022-03-07 16:41:20Z saprykin $
 * ===========================================================================
 *
 *                            PUBLIC DOMAIN NOTICE
 *               National Center for Biotechnology Information
 *
 *  This software/database is a "United States Government Work" under the
 *  terms of the United States Copyright Act.  It was written as part of
 *  the author's official duties as a United States Government employee and
 *  thus cannot be copyrighted.  This software/database is freely available
 *  to the public for use. The National Library of Medicine and the U.S.
 *  Government have not placed any restriction on its use or reproduction.
 *
 *  Although all reasonable efforts have been taken to ensure the accuracy
 *  and reliability of the software and data, the NLM and the U.S.
 *  Government do not and cannot warrant the performance or results that
 *  may be obtained by using this software or data. The NLM and the U.S.
 *  Government disclaim all warranties, express or implied, including
 *  warranties of performance, merchantability or fitness for any particular
 *  purpose.
 *
 *  Please cite the author in any work or product based on this material.
 *
 *  Db Cassandra: class generating execution plans for cassandra table segmented scans.
 *
 */

#include <ncbi_pch.hpp>

#include <objtools/pubseq_gateway/impl/cassandra/fullscan/seg_plan.hpp>

BEGIN_IDBLOB_SCOPE
USING_NCBI_SCOPE;

BEGIN_SCOPE()
void FilterTokenRanges(CCassConnection::TTokenRanges& ranges, size_t segment, size_t segment_count)
{
    if (segment_count > 1) {
        for (ssize_t i = ranges.size() - 1; i >= 0; --i) {
            if (i % segment_count != segment) {
                ranges.erase(ranges.begin() + i);
            }
        }
    }
}
END_SCOPE()

CCassandraSegscanPlan::CCassandraSegscanPlan()
  : m_Segment({0, 1})
{}

CCassandraSegscanPlan& CCassandraSegscanPlan::SetSegment(pair<size_t, size_t> segment)
{
    m_Segment = segment;
    return *this;
}

void CCassandraSegscanPlan::Generate()
{
    auto partition_count_limit = GetPartitionCountPerQueryLimit();
    if (partition_count_limit > 0) {
        // Limiting token range partition size is unstable as size estimates may be different on different nodes
        //  so we first apply segment filter
        //  and then we split remaining token ranges accorting to limit
        SetPartitionCountPerQueryLimit(0);
        CCassandraFullscanPlan::Generate();
        FilterTokenRanges(GetTokenRanges(), m_Segment.first, m_Segment.second);
        SetPartitionCountPerQueryLimit(partition_count_limit);
        SplitTokenRangesForLimits();
    }
    else {
        CCassandraFullscanPlan::Generate();
        FilterTokenRanges(GetTokenRanges(), m_Segment.first, m_Segment.second);
    }
}

END_IDBLOB_SCOPE
