co.cask.cdap.etl.common.PipelinePhase.java Source code

Java tutorial

Introduction

Here is the source code for co.cask.cdap.etl.common.PipelinePhase.java

Source

/*
 * Copyright  2016 Cask Data, Inc.
 *
 * 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 co.cask.cdap.etl.common;

import co.cask.cdap.etl.planner.Dag;
import co.cask.cdap.etl.planner.StageInfo;
import co.cask.cdap.etl.proto.Connection;
import com.google.common.base.Joiner;
import com.google.common.base.Predicate;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterators;
import com.google.common.collect.Sets;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;

/**
 * Keeps track of the plugin ids for the source, transforms, and sink of a pipeline phase.
 */
public class PipelinePhase implements Iterable<StageInfo> {
    // plugin type -> stage info
    private final Map<String, Set<StageInfo>> stages;
    private final Dag dag;

    private PipelinePhase(Map<String, Set<StageInfo>> stages, Dag dag) {
        this.stages = ImmutableMap.copyOf(stages);
        this.dag = dag;
    }

    /**
     * Get an unmodifiable set of stages that use the specified plugin type.
     *
     * @param pluginType the plugin type
     * @return unmodifiable set of stages that use the specified plugin type
     */
    public Set<StageInfo> getStagesOfType(String pluginType) {
        Set<StageInfo> stageInfos = stages.get(pluginType);
        return Collections.unmodifiableSet(stageInfos == null ? new HashSet<StageInfo>() : stageInfos);
    }

    public Set<String> getStageOutputs(String stage) {
        Set<String> outputs = dag.getNodeOutputs(stage);
        return Collections.unmodifiableSet(outputs == null ? new HashSet<String>() : outputs);
    }

    public Set<String> getPluginTypes() {
        return stages.keySet();
    }

    public Set<String> getSources() {
        return dag.getSources();
    }

    public Set<String> getSinks() {
        return dag.getSinks();
    }

    /**
     * Get a subset of the pipeline phase, starting from the sources and going to the specified nodes that will
     * be the new sinks of the pipeline subset.
     *
     * @param newSinks the new sinks to go to
     * @return subset of the pipeline, starting from current sources and going to the new sinks
     */
    public PipelinePhase subsetTo(Set<String> newSinks) {
        return getSubset(dag.subsetFrom(dag.getSources(), newSinks));
    }

    /**
     * Get a subset of the pipeline phase, starting from the specified new sources and going to the current sinks.
     *
     * @param newSources the new sources to start from
     * @return subset of the pipeline, starting from specified new sources and going to the current sinks
     */
    public PipelinePhase subsetFrom(Set<String> newSources) {
        return getSubset(dag.subsetFrom(newSources));
    }

    private PipelinePhase getSubset(final Dag subsetDag) {
        Map<String, Set<StageInfo>> subsetStages = new HashMap<>();
        for (Map.Entry<String, Set<StageInfo>> stagesEntry : stages.entrySet()) {
            final Set<StageInfo> stagesOfType = Sets.filter(stagesEntry.getValue(), new Predicate<StageInfo>() {
                @Override
                public boolean apply(StageInfo stageInfo) {
                    return subsetDag.getNodes().contains(stageInfo.getName());
                }
            });
            if (!stagesOfType.isEmpty()) {
                subsetStages.put(stagesEntry.getKey(), stagesOfType);
            }
        }
        return new PipelinePhase(subsetStages, subsetDag);
    }

    @Override
    public boolean equals(Object o) {
        if (this == o) {
            return true;
        }
        if (o == null || getClass() != o.getClass()) {
            return false;
        }

        PipelinePhase that = (PipelinePhase) o;

        return Objects.equals(stages, that.stages) && Objects.equals(dag, that.dag);
    }

    @Override
    public int hashCode() {
        return Objects.hash(stages, dag);
    }

    @Override
    public String toString() {
        return "PipelinePhase{" + "stages=" + stages + ", dag=" + dag + '}';
    }

    /**
     * Get a builder used to create a pipeline phase.
     *
     * @param supportedPluginTypes types of plugins supported in the phase
     * @return builder used to create a pipeline phase
     */
    public static Builder builder(Set<String> supportedPluginTypes) {
        return new Builder(supportedPluginTypes);
    }

    @Override
    public Iterator<StageInfo> iterator() {
        List<Iterator<StageInfo>> iterators = new ArrayList<>(stages.size());
        for (Map.Entry<String, Set<StageInfo>> stagesEntry : stages.entrySet()) {
            iterators.add(stagesEntry.getValue().iterator());
        }
        return Iterators.concat(iterators.iterator());
    }

    /**
     * Builder to create a {@link PipelinePhase}.
     */
    public static class Builder {
        private final Set<String> supportedPluginTypes;
        private final Map<String, Set<StageInfo>> stages;
        private final Set<co.cask.cdap.etl.proto.Connection> connections;

        public Builder(Set<String> supportedPluginTypes) {
            this.supportedPluginTypes = supportedPluginTypes;
            this.stages = new HashMap<>();
            this.connections = new HashSet<>();
        }

        public Builder addStage(String pluginType, StageInfo stageInfo) {
            return addStages(pluginType, ImmutableSet.of(stageInfo));
        }

        public Builder addStages(String pluginType, Collection<StageInfo> stages) {
            if (!supportedPluginTypes.contains(pluginType)) {
                throw new IllegalArgumentException(
                        String.format("%s is an unsupported plugin type. Plugin type must be one of %s.",
                                pluginType, Joiner.on(',').join(supportedPluginTypes)));
            }
            Set<StageInfo> existingStages = this.stages.get(pluginType);
            if (existingStages == null) {
                existingStages = new HashSet<>();
                this.stages.put(pluginType, existingStages);
            }
            existingStages.addAll(stages);
            return this;
        }

        public Builder addConnection(String from, String to) {
            return addConnections(from, ImmutableSet.of(to));
        }

        public Builder addConnections(String from, Collection<String> to) {
            for (String toStage : to) {
                connections.add(new Connection(from, toStage));
            }
            return this;
        }

        public Builder addConnections(Map<String, Set<String>> connections) {
            for (Map.Entry<String, Set<String>> entry : connections.entrySet()) {
                addConnections(entry.getKey(), entry.getValue());
            }
            return this;
        }

        public PipelinePhase build() {
            return new PipelinePhase(stages, new Dag(connections));
        }
    }
}