com.facebook.presto.jdbc.internal.spi.plan.SetOperationNode Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of presto-jdbc Show documentation
The newest version!
/*
 * Licensed 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 com.facebook.presto.jdbc.internal.spi.plan;

import com.facebook.presto.jdbc.internal.spi.SourceLocation;
import com.facebook.presto.jdbc.internal.spi.relation.VariableReferenceExpression;
import com.facebook.presto.jdbc.internal.jackson.annotation.JsonCreator;
import com.facebook.presto.jdbc.internal.jackson.annotation.JsonProperty;

import com.facebook.presto.jdbc.internal.javax.annotation.concurrent.Immutable;

import java.util.ArrayList;
import java.util.Collection;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;

import static java.lang.String.format;
import static java.util.Collections.unmodifiableList;
import static java.util.Collections.unmodifiableMap;
import static java.util.Objects.requireNonNull;
import static java.util.stream.Collectors.toList;

@Immutable
public abstract class SetOperationNode
        extends PlanNode
{
    private final List sources;
    private final Map> outputToInputs;
    private final List outputVariables;

    @JsonCreator
    protected SetOperationNode(
            Optional sourceLocation,
            @JsonProperty("id") PlanNodeId id,
            @JsonProperty("sources") List sources,
            @JsonProperty("outputVariables") List outputVariables,
            @JsonProperty("outputToInputs") Map> outputToInputs)
    {
        this(sourceLocation, id, Optional.empty(), sources, outputVariables, outputToInputs);
    }

    protected SetOperationNode(
            Optional sourceLocation,
            PlanNodeId id,
            Optional statsEquivalentPlanNode,
            List sources,
            List outputVariables,
            Map> outputToInputs)
    {
        super(sourceLocation, id, statsEquivalentPlanNode);

        requireNonNull(sources, "sources is null");
        checkArgument(!sources.isEmpty(), "Must have at least one source");
        requireNonNull(outputToInputs, "outputToInputs is null");

        this.sources = unmodifiableList(new ArrayList<>(sources));
        Map> copiedMap = new LinkedHashMap<>();
        outputToInputs.forEach((key, value) -> copiedMap.put(key, unmodifiableList(new ArrayList<>(value))));
        this.outputToInputs = unmodifiableMap(copiedMap);
        this.outputVariables = unmodifiableList(new ArrayList<>(outputVariables));

        for (Collection inputs : this.outputToInputs.values()) {
            checkArgument(
                    inputs.size() == this.sources.size(),
                    format("Every source needs to map its variables to an output %s operation variables", this.getClass().getSimpleName()));
        }

        // Make sure each source positionally corresponds to their variable values in the Multimap
        for (int i = 0; i < sources.size(); i++) {
            for (List expectedInputs : this.outputToInputs.values()) {
                checkArgument(sources.get(i).getOutputVariables().contains(expectedInputs.get(i)), "Source does not provide required variables");
            }
        }
    }

    @Override
    @JsonProperty
    public List getSources()
    {
        return sources;
    }

    @JsonProperty
    public Map> getVariableMapping()
    {
        return outputToInputs;
    }

    @Override
    @JsonProperty
    public List getOutputVariables()
    {
        return outputVariables;
    }

    public List sourceOutputLayout(int sourceIndex)
    {
        // Make sure the sourceOutputLayout variables are listed in the same order as the corresponding output variables
        return unmodifiableList(getOutputVariables().stream().map(variable -> outputToInputs.get(variable).get(sourceIndex)).collect(toList()));
    }

    /**
     * Returns the output to input variable mapping for the given source channel
     */
    public Map sourceVariableMap(int sourceIndex)
    {
        Map result = new LinkedHashMap<>();
        for (Map.Entry> entry : outputToInputs.entrySet()) {
            result.put(entry.getKey(), entry.getValue().get(sourceIndex));
        }

        return unmodifiableMap(result);
    }

    private static void checkArgument(boolean condition, String message)
    {
        if (!condition) {
            throw new IllegalArgumentException(message);
        }
    }
}