MTMVPartitionExpander.java

// 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.

package org.apache.doris.mtmv;

import org.apache.doris.catalog.PartitionItem;
import org.apache.doris.catalog.PartitionKey;
import org.apache.doris.catalog.PartitionType;
import org.apache.doris.catalog.RangePartitionItem;
import org.apache.doris.common.AnalysisException;
import org.apache.doris.datasource.mvcc.MvccSnapshot;
import org.apache.doris.datasource.mvcc.MvccUtil;

import com.google.common.collect.Maps;
import com.google.common.collect.Range;
import com.google.common.collect.Sets;

import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.NavigableMap;
import java.util.Optional;
import java.util.Set;
import java.util.TreeMap;

/**
 * Utility to expand query-used partition filters to MV partition granularity
 * using Range.encloses(), avoiding expensive dateTrunc / strToDate / dateIncrement
 * per-partition operations in the rollup pipeline.
 * Separated from MTMV to keep a lightweight dependency tree for testability ���
 * loading this class does not trigger MTMV ��� OlapTable ��� CloudReplica class loading.
 */
public class MTMVPartitionExpander {

    /**
     * Expand queryUsedPartitions to MV partition granularity for RANGE base tables.
     * For example, if MV is monthly partitioned (date_trunc(month)) and base table is daily:
     * - Query uses p_20250115 (Jan 15)
     * - Find MV partition p_202501 that encloses [20250115, 20250116)
     * - Expand to ALL daily partitions within p_202501's range [20250101, 20250201)
     * - Result: {p_20250101, p_20250102, ..., p_20250131}
     */
    public static Map<List<String>, Set<String>> expandToMvPartitionGranularity(
            Map<List<String>, Set<String>> queryUsedBaseTablePartitionMap,
            Map<String, PartitionItem> mvPartitionItems,
            Set<MTMVRelatedTableIf> pctTables) throws AnalysisException {
        NavigableMap<PartitionKey, Range<PartitionKey>> mvRanges = new TreeMap<>();
        for (PartitionItem item : mvPartitionItems.values()) {
            Range<PartitionKey> range = ((RangePartitionItem) item).getItems();
            mvRanges.put(range.lowerEndpoint(), range);
        }

        Map<List<String>, Set<String>> expanded = Maps.newHashMap();
        for (MTMVRelatedTableIf pctTable : pctTables) {
            List<String> qualifiers = pctTable.getFullQualifiers();
            Set<String> queryUsedPartitions = queryUsedBaseTablePartitionMap.get(qualifiers);
            if (queryUsedPartitions == null) {
                continue;
            }

            Optional<MvccSnapshot> snapshot = MvccUtil.getSnapshotFromContext(pctTable);
            if (pctTable.getPartitionType(snapshot) != PartitionType.RANGE) {
                expanded.put(qualifiers, queryUsedPartitions);
                continue;
            }

            Map<String, PartitionItem> basePartitionItems = pctTable.getAndCopyPartitionItems(snapshot);

            NavigableMap<PartitionKey, Range<PartitionKey>> relevantMvRanges = new TreeMap<>();
            for (String queriedBasePartition : queryUsedPartitions) {
                PartitionItem baseItem = basePartitionItems.get(queriedBasePartition);
                if (baseItem == null) {
                    continue;
                }
                Range<PartitionKey> baseRange = ((RangePartitionItem) baseItem).getItems();
                Range<PartitionKey> mvRange = findEnclosingRange(mvRanges, baseRange);
                if (mvRange != null) {
                    relevantMvRanges.put(mvRange.lowerEndpoint(), mvRange);
                }
            }

            if (relevantMvRanges.isEmpty()) {
                expanded.put(qualifiers, Sets.newHashSet());
                continue;
            }

            Set<String> expandedPartitions = Sets.newHashSet();
            for (Entry<String, PartitionItem> baseEntry : basePartitionItems.entrySet()) {
                Range<PartitionKey> baseRange = ((RangePartitionItem) baseEntry.getValue()).getItems();
                if (findEnclosingRange(relevantMvRanges, baseRange) != null) {
                    expandedPartitions.add(baseEntry.getKey());
                }
            }

            expanded.put(qualifiers, expandedPartitions);
        }

        return expanded;
    }

    private static Range<PartitionKey> findEnclosingRange(
            NavigableMap<PartitionKey, Range<PartitionKey>> ranges, Range<PartitionKey> baseRange) {
        // RANGE partitions do not overlap, so only the range with the nearest lower endpoint can enclose baseRange.
        Entry<PartitionKey, Range<PartitionKey>> candidate = ranges.floorEntry(baseRange.lowerEndpoint());
        return candidate != null && candidate.getValue().encloses(baseRange) ? candidate.getValue() : null;
    }

    private MTMVPartitionExpander() {
    }
}