双数据源混排

import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.tuple.Pair;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.function.Function;

/**
 * @author lanan
 */
public class PageQueryHelper <R, U> {

    private static final int MAX_QUERY_TIMES = 10;

    private static final int MAX_SORT_DATA_SIZE = 120;

    private final Function<PageQueryParam, Pair<Integer, List<R>>> aQuery;

    private final Function<PageQueryParam, Pair<Integer, List<R>>> bQuery;

    private final Function<R, U> compareFunction;

    private final Comparator<U> comparator;

    public PageQueryHelper(
            Function<PageQueryParam, Pair<Integer, List<R>>> aQuery,
            Function<PageQueryParam, Pair<Integer, List<R>>> bQuery,
            Function<R, U> compareFunction,
            final Comparator<U> comparator) {
        this.aQuery = aQuery;
        this.bQuery = bQuery;
        this.compareFunction = compareFunction;
        this.comparator = comparator;
    }


    /**
     * 双数据源混合查询
     * @param globalPageParam
     * @return
     */
    @SuppressWarnings({"PMD.NcssCount", "java:S3776"})
    public Pair<Integer, List<R>> dualDataSourcePageQuery(PageQueryParam globalPageParam) {
        if (globalPageParam == null) {
            return Pair.of(0, Collections.emptyList());
        }

        //分别取第一页数据，并将总页数一起取到
        final PageQueryParam aPageParam = initPageParam(globalPageParam);
        final PageQueryParam bPageParam = initPageParam(globalPageParam);
        final Pair<Integer, List<R>> aStartPage = aQuery.apply(aPageParam);
        final Pair<Integer, List<R>> bStartPage = bQuery.apply(bPageParam);
        final Integer aTotal = aStartPage.getLeft();
        final Integer bTotal = bStartPage.getLeft();

        final Integer globalStart = globalPageParam.getStart();
        //两个数据源的总数据量小于start
        final int total = aTotal + bTotal;
        if (total <= globalStart) {
            return Pair.of(total, Collections.emptyList());
        }

        //定义双数据源的取值范围（globalStart可能存在的数据范围）
        int aRight = Math.min(Math.max(0, aTotal - 1), globalStart);
        int bRight = Math.min(Math.max(0, bTotal - 1 ), globalStart);
        int aLeft = Math.max(globalStart - bTotal, 0);
        int bLeft = Math.max(globalStart - aTotal, 0);
        final Integer pageSize = globalPageParam.getNum();

        final List<R> aLeftPage = aLeft == 0 ? aStartPage.getRight() : queryA(aPageParam, aLeft, Math.min(pageSize, aTotal - aLeft));
        final List<R> bLeftPage = bLeft == 0 ? bStartPage.getRight() : queryB(bPageParam, bLeft, Math.min(pageSize, bTotal - bLeft));
        //优先处理这种情况：globalStart之前全是数据源a的数据或全是数据源b的数据
        final List<R> rs = returnIfSorted(aLeftPage, bLeftPage, aRight, bRight, aPageParam, bPageParam, globalPageParam);
        if (CollectionUtils.isNotEmpty(rs)) {
            return Pair.of(total, rs);
        }
        int times = 0;
        while (aRight > aLeft && ++times < MAX_QUERY_TIMES) {
            int aMid = (aRight - aLeft) / 2 + aLeft;
            int bMid = globalStart - aMid;
            // 如果 aRight - aLeft <= pageSize，则page页数据的取值访问为：
            // A[aLeft - pageSize, aRight + pageSize] + B[bMid - pageSize, bMid + 2 * pageSize]
            // 直接取到这个区间的数据，进行排序后，返回globalStart后的数据即可
            if (aRight - aLeft <= pageSize || bRight - bLeft <= pageSize) {
                if (aRight - aLeft <= pageSize) {
                    final List<R> result = new ArrayList<>(MAX_SORT_DATA_SIZE);
                    result.addAll(queryA(aPageParam, aLeft - pageSize, aRight - aLeft + 2 * pageSize));
                    result.addAll(queryB(bPageParam, bMid - pageSize,  3 * pageSize));
                    sort(result);
                    int skip = globalStart - Math.max(aLeft - pageSize, 0) - Math.max(bMid - pageSize, 0);
                    return Pair.of(total, result.subList(skip, Math.min(skip + pageSize, result.size())));
                } else {
                    final List<R> result = new ArrayList<>(MAX_SORT_DATA_SIZE);
                    result.addAll(queryB(bPageParam, bLeft - pageSize, bRight - bLeft + 2 * pageSize));
                    result.addAll(queryA(aPageParam, aMid - pageSize,   3 * pageSize));
                    sort(result);
                    int skip = globalStart - Math.max(bLeft - pageSize, 0) - Math.max(aMid - pageSize, 0);
                    return Pair.of(total, result.subList(skip, Math.min(skip + pageSize, result.size())));
                }
            }
            final List<R> aPage = queryA(aPageParam, aMid, pageSize);
            final List<R> bPage = queryB(bPageParam, bMid, pageSize);
            R aMin = aPage.get(0);
            R aMax = aPage.get(aPage.size() - 1);
            R bMin = bPage.get(0);
            R bMax = bPage.get(bPage.size() - 1);
            if (compare(aMin, bMax) > 0) {
                aRight = aMid - 1;
                bLeft = bMid;
            } else if (compare(bMin, aMax) > 0) {
                aLeft = aMid + 1;
                bRight = bMid;
            } else {
                //定位到了全局排序值大概位置
                //此时第page页的数据范围：
                //A[aMid - pageSize, aMid + pageSize]
                //B[bMid - pageSize, bMid + pageSize]
                //全部取回后排序，返回globalStart之后的数据即可
                final List<R> result = new ArrayList<>(MAX_SORT_DATA_SIZE);
                result.addAll(aPage);
                result.addAll(bPage);
                result.addAll(queryA(aPageParam, Math.max(aMid - pageSize, 0), Math.min(aMid, pageSize)));
                result.addAll(queryB(bPageParam, Math.max(bMid - pageSize, 0), Math.min(bMid, pageSize)));
                sort(result);
                int skip = globalStart - Math.max(aMid - pageSize, 0) - (Math.max(bMid - pageSize, 0));
                return Pair.of(total, result.subList(skip, pageSize + skip));
            }
        }

        LOG.warn("双数据混合查询异常");
        return Pair.of(total, Collections.emptyList());
    }

    /**
     * 特殊情况的处理，page页之前全是A或者全是B的数据，那么此时可直接返回需要的数据
     * aLeft + b
     * @param aLeftPage
     * @param bLeftPage
     * @param aRight
     * @param bRight
     * @param aPageParam
     * @param bPageParam
     * @param globalPageParam
     * @return
     */
    private List<R> returnIfSorted(List<R> aLeftPage, List<R> bLeftPage,
                                   int aRight, int bRight, PageQueryParam aPageParam, PageQueryParam bPageParam,
                                   PageQueryParam globalPageParam) {
        int pageSize = globalPageParam.getNum();
        //优先处理这种情况：globalStart之前全是数据源a的数据或全是数据源b的数据
        final List<R> bRightPage = queryB(bPageParam, bRight, pageSize);
        if (CollectionUtils.isEmpty(aLeftPage)) {
            return bRightPage;
        }
        int globalStart = globalPageParam.getStart();
        if (CollectionUtils.isEmpty(bRightPage) || compare(aLeftPage.get(0), bRightPage.get(0)) >= 0) {
            //bRight < globalStart，说明此时 bRightPage取的是 B数据源的最后一条数据，且该数据小于等于aLeftPage[0]
            if (bRight < globalStart) {
                return aLeftPage;
            } else {
                //bRightPage[0]之前的数据全小于aLeftPage[0]，此时第page页的数据在bRightPage+aLeftPage中
                final List<R> result = new ArrayList<>(aLeftPage.size() + bRightPage.size());
                result.addAll(aLeftPage);
                result.addAll(bRightPage);
                sort(result);
                return result.subList(0, Math.min(result.size(), pageSize));
            }
        }
        final List<R> aRightPage = queryA(aPageParam, aRight, pageSize);
        if (CollectionUtils.isEmpty(bLeftPage)) {
            return aRightPage;
        }
        if (CollectionUtils.isEmpty(aRightPage) || compare(bLeftPage.get(0), aRightPage.get(0)) >= 0) {
            if (aRight <  globalStart) {
                return bLeftPage;
            } else {
                final List<R> result = new ArrayList<>(aRightPage.size() + bLeftPage.size());
                result.addAll(aRightPage);
                result.addAll(bLeftPage);
                sort(result);
                return result.subList(0, Math.min(result.size(), pageSize));
            }
        }
        return Collections.emptyList();
    }

    /**
     * 比较两个数据
     * @param o1
     * @param o2
     * @return
     */
    private int compare(R o1, R o2) {
        final U aCompareValue = compareFunction.apply(o1);
        final U bCompareValue = compareFunction.apply(o2);
        return comparator.compare(aCompareValue, bCompareValue);
    }

    /**
     * 根据指定字段进行排序
     * */
    private void sort(List<R> list) {
        Collections.sort(list, this::compare);
    }

    /**
     * 生成查询条件
     * @param globalPageParam
     * @return
     */
    private PageQueryParam initPageParam(PageQueryParam globalPageParam) {
        final PageQueryParam pageParam = new PageQueryParam();
        pageParam.setUserId(globalPageParam.getUserId());
        pageParam.setKeyword(globalPageParam.getKeyword());
        pageParam.setStart(0);
        pageParam.setNum(globalPageParam.getNum());
        return pageParam;
    }

    /**
     * 从A数据源查询数据
     * @param pageParam 减少PageQueryParam对象创建数
     * @param start 起始范围
     * @param num 获取数量
     * @return
     */
    private List<R> queryA(PageQueryParam pageParam, int start, int num) {
        pageParam.setStart(start);
        pageParam.setNum(num);
        return queryList(aQuery, pageParam);
    }

    /**
     * 从B数据源查询数据
     * @param pageParam 减少PageQueryParam对象创建数
     * @param start 起始范围
     * @param num 获取数量
     * @return
     */
    private List<R> queryB(PageQueryParam pageParam, int start, int num) {
        pageParam.setStart(start);
        pageParam.setNum(num);
        return queryList(bQuery, pageParam);
    }

    /**
     * 从数据源获取数据
     * @param query
     * @param pageParam
     * @return
     */
    private List<R> queryList(Function<PageQueryParam, Pair<Integer, List<R>>> query, PageQueryParam pageParam) {
        if (pageParam.getStart() < 0) {
            pageParam.setStart(0);
        }
        final Pair<Integer, List<R>> pair = query.apply(pageParam);
        if (pair == null || CollectionUtils.isEmpty(pair.getRight())) {
            return Collections.emptyList();
        }
        return pair.getRight();
    }
}