[Java Streaming] Fix StreamSource constructor and SourceFunction initialization (#4597)

This commit is contained in:
ppeagle
2019-04-10 20:02:11 +08:00
committed by Hao Chen
parent ed02bf11f7
commit 6697407ec4
2 changed files with 5 additions and 3 deletions
@@ -2,6 +2,7 @@ package org.ray.streaming.api.stream;
import java.util.Collection;
import org.ray.streaming.api.context.StreamingContext;
import org.ray.streaming.api.function.impl.SourceFunction;
import org.ray.streaming.api.function.internal.CollectionSourceFunction;
import org.ray.streaming.operator.impl.SourceOperator;
@@ -12,8 +13,8 @@ import org.ray.streaming.operator.impl.SourceOperator;
*/
public class StreamSource<T> extends DataStream<T> {
public StreamSource(StreamingContext streamingContext, SourceOperator sourceOperator) {
super(streamingContext, sourceOperator);
public StreamSource(StreamingContext streamingContext, SourceFunction<T> sourceFunction) {
super(streamingContext, new SourceOperator<>(sourceFunction));
}
/**
@@ -25,7 +26,7 @@ public class StreamSource<T> extends DataStream<T> {
* @return A StreamSource.
*/
public static <T> StreamSource<T> buildSource(StreamingContext context, Collection<T> values) {
return new StreamSource(context, new SourceOperator(new CollectionSourceFunction(values)));
return new StreamSource(context, new CollectionSourceFunction(values));
}
public StreamSource<T> setParallelism(int parallelism) {
@@ -21,6 +21,7 @@ public class SourceOperator<T> extends StreamOperator<SourceFunction<T>> {
public void open(List<Collector> collectorList, RuntimeContext runtimeContext) {
super.open(collectorList, runtimeContext);
this.sourceContext = new SourceContextImpl(collectorList);
this.function.init(runtimeContext.getParallelism(), runtimeContext.getTaskIndex());
}
public void process(Long batchId) {