MatchInputResolver.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.nereids.rules.analysis;

import org.apache.doris.nereids.CascadesContext;
import org.apache.doris.nereids.trees.expressions.Alias;
import org.apache.doris.nereids.trees.expressions.Cast;
import org.apache.doris.nereids.trees.expressions.Expression;
import org.apache.doris.nereids.trees.expressions.Slot;
import org.apache.doris.nereids.trees.expressions.SlotReference;
import org.apache.doris.nereids.trees.plans.Plan;
import org.apache.doris.nereids.trees.plans.logical.LogicalCTEConsumer;
import org.apache.doris.nereids.trees.plans.logical.LogicalCTEProducer;
import org.apache.doris.nereids.trees.plans.logical.LogicalGenerate;
import org.apache.doris.nereids.trees.plans.logical.LogicalRecursiveUnion;
import org.apache.doris.nereids.trees.plans.logical.LogicalSetOperation;
import org.apache.doris.nereids.trees.plans.logical.LogicalWorkTableReference;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

/** Resolves only MATCH input semantics; it does not attach lineage metadata to slots or plans. */
public final class MatchInputResolver {
    /** Proven scalar, possible values root, or a producer that is not yet resolved. */
    public enum InputKind { SCALAR, ROOT_CANDIDATE, UNRESOLVED }

    private final CascadesContext cascadesContext;
    private final Map<Slot, InputKind> cteOutputs = new HashMap<>();

    public MatchInputResolver(CascadesContext cascadesContext) {
        this.cascadesContext = cascadesContext;
    }

    /** Resolve aliases and positional producer mappings without choosing an analyzer. */
    public InputKind resolve(Expression input, Plan plan) {
        // Optimizer VARIANT layout casts preserve root identity. SQL casts have already received
        // their scalar STRING coercion during analysis, before expression replacement can occur.
        if (input instanceof Cast && input.getDataType().isVariantType()
                && input.child(0).getDataType().isVariantType()) {
            return resolve(input.child(0), plan);
        }
        if (!(input instanceof SlotReference)) {
            return InputKind.SCALAR;
        }
        SlotReference slot = (SlotReference) input;
        if (slot.hasSubColPath()) {
            return InputKind.SCALAR;
        }
        if (plan == null) {
            return slot.getOriginalColumn().isPresent() ? InputKind.ROOT_CANDIDATE : InputKind.UNRESOLVED;
        }
        if (plan instanceof LogicalCTEConsumer) {
            LogicalCTEConsumer consumer = (LogicalCTEConsumer) plan;
            Slot producerSlot = consumer.getProducerSlot(slot);
            InputKind cached = cteOutputs.get(producerSlot);
            if (cached != null) {
                return cached;
            }
            LogicalCTEProducer<?> producer = cascadesContext.getStatementContext()
                    .getCteProducerByCteId(consumer.getCteId());
            // Install a sentinel before walking a producer; unresolved cycles never become scalar.
            cteOutputs.put(producerSlot, InputKind.UNRESOLVED);
            InputKind result = producer == null ? InputKind.UNRESOLVED : resolve(producerSlot, producer.child());
            cteOutputs.put(producerSlot, result);
            return result;
        }
        if (plan instanceof LogicalSetOperation) {
            return resolveSetOutput(slot, plan, ((LogicalSetOperation) plan).getRegularChildrenOutputs());
        }
        if (plan instanceof LogicalRecursiveUnion) {
            return resolveSetOutput(slot, plan, ((LogicalRecursiveUnion<?, ?>) plan).getRegularChildrenOutputs());
        }
        // Generated values and materialized recursive worktable cells have scalar MATCH semantics;
        // neither denotes an independently bindable base-column values index.
        if (plan instanceof LogicalGenerate
                && ((LogicalGenerate<?>) plan).getGeneratorOutput().contains(slot)) {
            return InputKind.SCALAR;
        }
        if (plan instanceof LogicalWorkTableReference && plan.getOutputSet().contains(slot)) {
            return InputKind.SCALAR;
        }
        for (Expression expression : plan.getExpressions()) {
            if (expression instanceof Alias && ((Alias) expression).getExprId().equals(slot.getExprId())) {
                return resolve(expression.child(0), plan);
            }
        }
        for (Plan child : plan.children()) {
            if (child.getOutputSet().contains(slot)) {
                return resolve(slot, child);
            }
        }
        return slot.getOriginalColumn().isPresent() ? InputKind.ROOT_CANDIDATE : InputKind.UNRESOLVED;
    }

    private InputKind resolveSetOutput(Slot slot, Plan plan, List<List<SlotReference>> outputs) {
        int index = plan.getOutput().indexOf(slot);
        if (index < 0) {
            return InputKind.UNRESOLVED;
        }
        InputKind result = InputKind.SCALAR;
        for (int i = 0; i < plan.arity(); i++) {
            InputKind child = resolve(outputs.get(i).get(index), plan.child(i));
            if (child == InputKind.ROOT_CANDIDATE) {
                return child;
            }
            if (child == InputKind.UNRESOLVED) {
                result = child;
            }
        }
        return result;
    }
}