diff --git a/java/streaming/src/main/java/org/ray/streaming/api/stream/StreamSource.java b/java/streaming/src/main/java/org/ray/streaming/api/stream/StreamSource.java
index 727410d5a..05d1ddd91 100644
--- a/java/streaming/src/main/java/org/ray/streaming/api/stream/StreamSource.java
+++ b/java/streaming/src/main/java/org/ray/streaming/api/stream/StreamSource.java
@@ -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) {
diff --git a/java/streaming/src/main/java/org/ray/streaming/operator/impl/SourceOperator.java b/java/streaming/src/main/java/org/ray/streaming/operator/impl/SourceOperator.java
index 679000abd..430bb78c5 100644
--- a/java/streaming/src/main/java/org/ray/streaming/operator/impl/SourceOperator.java
+++ b/java/streaming/src/main/java/org/ray/streaming/operator/impl/SourceOperator.java
@@ -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) {