com.hazelcast.jet.impl.pipeline.PipelineImpl Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of hazelcast-jet Show documentation
There is a newer version: 4.5.4
/*
 * Copyright (c) 2008-2018, Hazelcast, Inc. All Rights Reserved.
 *
 * 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.hazelcast.jet.impl.pipeline;

import com.hazelcast.jet.core.DAG;
import com.hazelcast.jet.impl.pipeline.transform.BatchSourceTransform;
import com.hazelcast.jet.impl.pipeline.transform.SinkTransform;
import com.hazelcast.jet.impl.pipeline.transform.StreamSourceTransform;
import com.hazelcast.jet.impl.pipeline.transform.Transform;
import com.hazelcast.jet.pipeline.BatchSource;
import com.hazelcast.jet.pipeline.BatchStage;
import com.hazelcast.jet.pipeline.GeneralStage;
import com.hazelcast.jet.pipeline.Pipeline;
import com.hazelcast.jet.pipeline.Sink;
import com.hazelcast.jet.pipeline.SinkStage;
import com.hazelcast.jet.pipeline.StreamSource;
import com.hazelcast.jet.pipeline.StreamStage;

import javax.annotation.Nonnull;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.IntStream;

import static com.hazelcast.jet.impl.pipeline.ComputeStageImplBase.ADAPT_TO_JET_EVENT;
import static com.hazelcast.jet.impl.pipeline.ComputeStageImplBase.DONT_ADAPT;
import static java.util.stream.Collectors.toList;

public class PipelineImpl implements Pipeline {

    private final Map> adjacencyMap = new HashMap<>();

    @Nonnull @Override
    @SuppressWarnings("unchecked")
    public  BatchStage drawFrom(@Nonnull BatchSource source) {
        return new BatchStageImpl<>((BatchSourceTransform) source, this);
    }

    @Nonnull @Override
    @SuppressWarnings("unchecked")
    public  StreamStage drawFrom(@Nonnull StreamSource source) {
        StreamSourceTransform xform = (StreamSourceTransform) source;
        return new StreamStageImpl<>(xform, xform.emitsJetEvents() ? ADAPT_TO_JET_EVENT : DONT_ADAPT, this);
    }

    @Override
    public  SinkStage drainTo(@Nonnull Sink sink, GeneralStage... stagesToDrain) {
        if (stagesToDrain == null || stagesToDrain.length == 0) {
            throw new IllegalArgumentException("No stages supplied to Pipeline.drainTo()");
        }
        List upstream = Arrays.stream(stagesToDrain)
                                         .map(s -> (AbstractStage) s)
                                         .map(s -> s.transform)
                                         .collect(toList());
        int[] ordinalsToAdapt = IntStream
                .range(0, stagesToDrain.length)
                .filter(i -> ((ComputeStageImplBase) stagesToDrain[i]).fnAdapter == ADAPT_TO_JET_EVENT)
                .toArray();
        SinkImpl sinkImpl = (SinkImpl) sink;
        SinkTransform sinkTransform = new SinkTransform(sinkImpl, upstream, ordinalsToAdapt);
        SinkStageImpl sinkStage = new SinkStageImpl(sinkTransform, this);
        sinkImpl.onAssignToStage();
        connect(upstream, sinkTransform);
        return sinkStage;
    }

    @Nonnull @Override
    public DAG toDag() {
        return new Planner(this).createDag();
    }

    public void connect(Transform upstream, Transform downstream) {
        adjacencyMap.get(upstream).add(downstream);
    }

    public void connect(List upstream, Transform downstream) {
        upstream.forEach(u -> connect(u, downstream));
    }

    @Override
    public String toString() {
        return "Pipeline " + adjacencyMap;
    }

    Map> adjacencyMap() {
        Map> safeCopy = new HashMap<>();
        adjacencyMap.forEach((k, v) -> safeCopy.put(k, new ArrayList<>(v)));
        return safeCopy;
    }

    void register(Transform stage, List downstream) {
        List prev = adjacencyMap.put(stage, downstream);
        assert prev == null : "Double registering of a Stage with this Pipeline: " + stage;
    }
}