Streaming rich function (#8602)

streaming/java/streaming-api/src/main/java/io/ray/streaming/api/context/RuntimeContext.java

@@ -20,14 +20,20 @@ public interface RuntimeContext {
 
   int getParallelism();
 
+  /**
+   * @return config of current function
+   */
+  Map<String, String> getConfig();
+
+  /**
+   * @return config of the job
+   */
+  Map<String, String> getJobConfig();
+
   Long getCheckpointId();
 
   void setCheckpointId(long checkpointId);
 
-  Long getMaxBatch();
-
-  Map<String, String> getConfig();
-
   void setCurrentKey(Object key);
 
   KeyStateBackend getKeyStateBackend();

streaming/java/streaming-api/src/main/java/io/ray/streaming/api/function/RichFunction.java

@@ -0,0 +1,24 @@
+package io.ray.streaming.api.function;
+
+import io.ray.streaming.api.context.RuntimeContext;
+
+/**
+ * An interface for all user-defined functions to define the life cycle methods of the
+ * functions, and access the task context where the functions get executed.
+ */
+public interface RichFunction extends Function {
+
+  /**
+   * Initialization method for user function which called before the first call to the user
+   * function.
+   * @param runtimeContext runtime context
+   */
+  void open(RuntimeContext runtimeContext);
+
+  /**
+   * Tear-down method for the user function which called after the last call to
+   * the user function.
+   */
+  void close();
+
+}

streaming/java/streaming-api/src/main/java/io/ray/streaming/api/function/internal/Functions.java

@@ -0,0 +1,40 @@
+package io.ray.streaming.api.function.internal;
+
+import io.ray.streaming.api.context.RuntimeContext;
+import io.ray.streaming.api.function.Function;
+import io.ray.streaming.api.function.RichFunction;
+
+/**
+ * A util class for {@link Function}
+ */
+public class Functions {
+
+  private static class DefaultRichFunction implements RichFunction {
+    private final Function function;
+
+    private DefaultRichFunction(Function function) {
+      this.function = function;
+    }
+
+    @Override
+    public void open(RuntimeContext runtimeContext) {
+    }
+
+    @Override
+    public void close() {
+    }
+
+    public Function getFunction() {
+      return function;
+    }
+  }
+
+  public static RichFunction wrap(Function function) {
+    if (function instanceof RichFunction) {
+      return (RichFunction) function;
+    } else {
+      return new DefaultRichFunction(function);
+    }
+  }
+
+}

streaming/java/streaming-api/src/main/java/io/ray/streaming/api/stream/Stream.java

@@ -9,6 +9,8 @@ import io.ray.streaming.operator.Operator;
 import io.ray.streaming.operator.StreamOperator;
 import io.ray.streaming.python.PythonPartition;
 import java.io.Serializable;
+import java.util.HashMap;
+import java.util.Map;
 
 /**
  * Abstract base class of all stream types.
@@ -23,6 +25,7 @@ public abstract class Stream<S extends Stream<S, T>, T>
   private final Stream inputStream;
   private final StreamOperator operator;
   private int parallelism = 1;
+  private Map<String, String> config = new HashMap<>();
   private Partition<T> partition;
   private Stream originalStream;
 
@@ -134,6 +137,25 @@ public abstract class Stream<S extends Stream<S, T>, T>
     return self();
   }
 
+  public S withConfig(Map<String, String> config) {
+    config.forEach(this::withConfig);
+    return self();
+  }
+
+  public S withConfig(String key, String value) {
+    if (isProxyStream()) {
+      originalStream.withConfig(key, value);
+    } else {
+      this.config.put(key, value);
+    }
+    return self();
+  }
+
+  @SuppressWarnings("unchecked")
+  public Map<String, String> getConfig() {
+    return isProxyStream() ? originalStream.getConfig() : config;
+  }
+
   public boolean isProxyStream() {
     return originalStream != null;
   }

streaming/java/streaming-api/src/main/java/io/ray/streaming/jobgraph/JobGraphBuilder.java

@@ -77,6 +77,7 @@ public class JobGraphBuilder {
     } else {
       throw new UnsupportedOperationException("Unsupported stream: " + stream);
     }
+    jobVertex.setConfig(stream.getConfig());
     this.jobGraph.addVertex(jobVertex);
   }
 

streaming/java/streaming-api/src/main/java/io/ray/streaming/jobgraph/JobVertex.java

@@ -4,20 +4,23 @@ import com.google.common.base.MoreObjects;
 import io.ray.streaming.api.Language;
 import io.ray.streaming.operator.StreamOperator;
 import java.io.Serializable;
+import java.util.Map;
 
 /**
  * Job vertex is a cell node where logic is executed.
  */
 public class JobVertex implements Serializable {
-
   private int vertexId;
   private int parallelism;
   private VertexType vertexType;
   private Language language;
   private StreamOperator streamOperator;
+  private Map<String, String> config;
 
-  public JobVertex(int vertexId, int parallelism, VertexType vertexType,
-      StreamOperator streamOperator) {
+  public JobVertex(int vertexId,
+                   int parallelism,
+                   VertexType vertexType,
+                   StreamOperator streamOperator) {
     this.vertexId = vertexId;
     this.parallelism = parallelism;
     this.vertexType = vertexType;
@@ -45,6 +48,14 @@ public class JobVertex implements Serializable {
     return language;
   }
 
+  public Map<String, String> getConfig() {
+    return config;
+  }
+
+  public void setConfig(Map<String, String> config) {
+    this.config = config;
+  }
+
   @Override
   public String toString() {
     return MoreObjects.toStringHelper(this)
@@ -53,7 +64,7 @@ public class JobVertex implements Serializable {
         .add("vertexType", vertexType)
         .add("language", language)
         .add("streamOperator", streamOperator)
+        .add("config", config)
         .toString();
   }
-
 }

streaming/java/streaming-api/src/main/java/io/ray/streaming/operator/StreamOperator.java

@@ -4,25 +4,30 @@ import io.ray.streaming.api.Language;
 import io.ray.streaming.api.collector.Collector;
 import io.ray.streaming.api.context.RuntimeContext;
 import io.ray.streaming.api.function.Function;
+import io.ray.streaming.api.function.RichFunction;
+import io.ray.streaming.api.function.internal.Functions;
 import io.ray.streaming.message.KeyRecord;
 import io.ray.streaming.message.Record;
 import java.util.List;
 
 public abstract class StreamOperator<F extends Function> implements Operator {
-  protected String name;
-  protected F function;
+  protected final String name;
+  protected final F function;
+  protected final RichFunction richFunction;
   protected List<Collector> collectorList;
   protected RuntimeContext runtimeContext;
 
   public StreamOperator(F function) {
     this.name = getClass().getSimpleName();
     this.function = function;
+    this.richFunction = Functions.wrap(function);
   }
 
   @Override
   public void open(List<Collector> collectorList, RuntimeContext runtimeContext) {
     this.collectorList = collectorList;
     this.runtimeContext = runtimeContext;
+    richFunction.open(runtimeContext);
   }
 
   @Override
@@ -32,7 +37,7 @@ public abstract class StreamOperator<F extends Function> implements Operator {
 
   @Override
   public void close() {
-
+    richFunction.close();
   }
 
   @Override

streaming/java/streaming-api/src/main/java/io/ray/streaming/util/Config.java

@@ -2,17 +2,6 @@ package io.ray.streaming.util;
 
 public class Config {
 
-  /**
-   * Maximum number of batches to run in a streaming job.
-   */
-  public static final String STREAMING_BATCH_MAX_COUNT = "streaming.batch.max.count";
-
-  /**
-   * batch frequency in milliseconds
-   */
-  public static final String STREAMING_BATCH_FREQUENCY = "streaming.batch.frequency";
-  public static final long STREAMING_BATCH_FREQUENCY_DEFAULT = 1000;
-
   public static final String STREAMING_JOB_NAME = "streaming.job.name";
   public static final String STREAMING_OP_NAME = "streaming.op_name";
   public static final String STREAMING_WORKER_NAME = "streaming.worker_name";