Examples with TBaseGraph edu.iu.dsc.tws.tset.TBaseGraph used on opensource projects

Example 1 with TBaseGraph

use of edu.iu.dsc.tws.tset.TBaseGraph in project beam by apache.

the class BeamBatchWorker method addInputs.

/**
 * Adds all the side inputs into the sink test so it is available from the DoFn's.
 */
private void addInputs(BaseTSet sinkTSet, Map<String, CachedTSet> sideInputTSets) {
    if (sideInputTSets.isEmpty()) {
        return;
    }
    TBaseGraph graph = sinkTSet.getTBaseGraph();
    TBase currNode = null;
    Deque<TBase> deque = new ArrayDeque<>();
    deque.add(sinkTSet);
    while (!deque.isEmpty()) {
        currNode = deque.remove();
        deque.addAll(graph.getPredecessors(currNode));
        if (currNode instanceof ComputeTSet) {
            if (((ComputeTSet) currNode).getComputeFunc() instanceof DoFnFunction) {
                Set<String> sideInputKeys = ((DoFnFunction) ((ComputeTSet) currNode).getComputeFunc()).getSideInputKeys();
                for (String sideInputKey : sideInputKeys) {
                    if (!sideInputTSets.containsKey(sideInputKey)) {
                        throw new IllegalStateException("Side input not found for key " + sideInputKey);
                    }
                    ((ComputeTSet) currNode).addInput(sideInputKey, sideInputTSets.get(sideInputKey));
                }
            }
        }
    }
}

Example 2 with TBaseGraph

use of edu.iu.dsc.tws.tset.TBaseGraph in project beam by apache.

the class BeamBatchWorker method execute.

@Override
public void execute(BatchTSetEnvironment env) {
    Config config = env.getConfig();
    Map<String, String> sideInputIds = (LinkedHashMap<String, String>) config.get(SIDEINPUTS);
    Set<String> leaveIds = (Set<String>) config.get(LEAVES);
    TBaseGraph graph = (TBaseGraph) config.get(GRAPH);
    env.settBaseGraph(graph);
    setupTSets(env, sideInputIds, leaveIds);
    resetEnv(env, graph);
    executePipeline(env);
}