[Streaming] Streaming filter transform (#6816)

* add filter transform

* lint

* add new line

streaming/java/streaming-api/src/main/java/org/ray/streaming/api/function/impl/FilterFunction.java

@@ -0,0 +1,21 @@
+package org.ray.streaming.api.function.impl;
+
+import org.ray.streaming.api.function.Function;
+
+/**
+ * A filter function is a predicate applied individually to each record.
+ * The predicate decides whether to keep the element, or to discard it.
+ *
+ * @param <T> type of the input data.
+ */
+@FunctionalInterface
+public interface FilterFunction<T> extends Function {
+
+  /**
+   * The filter function that evaluates the predicate.
+   *
+   * @param value The value to be filtered.
+   * @return True for values that should be retained, false for values to be filtered out.
+   */
+  boolean filter(T value) throws Exception;
+}

streaming/java/streaming-api/src/main/java/org/ray/streaming/api/stream/DataStream.java

@@ -2,6 +2,7 @@ package org.ray.streaming.api.stream;
 
 
 import org.ray.streaming.api.context.StreamingContext;
+import org.ray.streaming.api.function.impl.FilterFunction;
 import org.ray.streaming.api.function.impl.FlatMapFunction;
 import org.ray.streaming.api.function.impl.KeyFunction;
 import org.ray.streaming.api.function.impl.MapFunction;
@@ -9,6 +10,7 @@ import org.ray.streaming.api.function.impl.SinkFunction;
 import org.ray.streaming.api.partition.Partition;
 import org.ray.streaming.api.partition.impl.BroadcastPartition;
 import org.ray.streaming.operator.StreamOperator;
+import org.ray.streaming.operator.impl.FilterOperator;
 import org.ray.streaming.operator.impl.FlatMapOperator;
 import org.ray.streaming.operator.impl.KeyByOperator;
 import org.ray.streaming.operator.impl.MapOperator;
@@ -53,6 +55,10 @@ public class DataStream<T> extends Stream<T> {
     return new DataStream(this, new FlatMapOperator(flatMapFunction));
   }
 
+  public DataStream<T> filter(FilterFunction<T> filterFunction) {
+    return new DataStream<T>(this, new FilterOperator(filterFunction));
+  }
+
   /**
    * Apply a union transformation to this stream, with another stream.
    *

streaming/java/streaming-api/src/main/java/org/ray/streaming/jobgraph/JobGraph.java

@@ -4,6 +4,7 @@ import java.io.Serializable;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -30,17 +31,14 @@ public class JobGraph implements Serializable {
   /**
    * Generate direct-graph(made up of a set of vertices and connected by edges)
    * by current job graph for simple log printing.
-   * @return Digraph in string type.
    *
-   * Notice:
-   * This is temporarily implemented in hard code.
-   * May use 'guru.nidi:graphviz-java' as 3rd dependency in the future if needed.
+   * @return Digraph in string type.
    */
   public String generateDigraph() {
     StringBuilder digraph = new StringBuilder();
     digraph.append("digraph ").append(jobName + " ").append(" {");
 
-    for (JobEdge jobEdge: jobEdgeList) {
+    for (JobEdge jobEdge : jobEdgeList) {
       String srcNode = null;
       String targetNode = null;
       for (JobVertex jobVertex : jobVertexList) {

streaming/java/streaming-api/src/main/java/org/ray/streaming/operator/impl/FilterOperator.java

@@ -0,0 +1,21 @@
+package org.ray.streaming.operator.impl;
+
+import org.ray.streaming.api.function.impl.FilterFunction;
+import org.ray.streaming.message.Record;
+import org.ray.streaming.operator.OneInputOperator;
+import org.ray.streaming.operator.StreamOperator;
+
+public class FilterOperator<T> extends StreamOperator<FilterFunction<T>> implements
+    OneInputOperator<T> {
+
+  public FilterOperator(FilterFunction<T> filterFunction) {
+    super(filterFunction);
+  }
+
+  @Override
+  public void processElement(Record<T> record) throws Exception {
+    if (this.function.filter(record.getValue())) {
+      this.collect(record);
+    }
+  }
+}

streaming/java/streaming-api/src/main/java/org/ray/streaming/python/stream/PythonDataStream.java

@@ -43,6 +43,18 @@ public class PythonDataStream extends Stream implements PythonStream {
     return new PythonDataStream(this, new PythonOperator(func));
   }
 
+  /**
+   * Apply a filter function to this stream.
+   *
+   * @param func The python FilterFunction.
+   * @return A new PythonDataStream that contains only the elements satisfying
+   *     the given filter predicate.
+   */
+  public PythonDataStream filter(PythonFunction func) {
+    func.setFunctionInterface(FunctionInterface.FILTER_FUNCTION);
+    return new PythonDataStream(this, new PythonOperator(func));
+  }
+
   /**
    * Apply a sink function and get a StreamSink.
    *

streaming/java/streaming-runtime/src/test/java/org/ray/streaming/runtime/demo/WordCountTest.java

@@ -44,6 +44,7 @@ public class WordCountTest extends BaseUnitTest implements Serializable {
             collector.collect(new WordAndCount(record, 1));
           }
         })
+        .filter(pair -> !pair.word.contains("world"))
         .keyBy(pair -> pair.word)
         .reduce((ReduceFunction<WordAndCount>) (oldValue, newValue) ->
             new WordAndCount(oldValue.word, oldValue.count + newValue.count))
@@ -53,14 +54,14 @@ public class WordCountTest extends BaseUnitTest implements Serializable {
     streamingContext.execute("testWordCount");
 
     // Sleep until the count for every word is computed.
-    while (wordCount.size() < 3) {
+    while (wordCount.size() < 2) {
       try {
         Thread.sleep(100);
       } catch (InterruptedException e) {
         LOGGER.warn("Got an exception while sleeping.", e);
       }
     }
-    Assert.assertEquals(wordCount, ImmutableMap.of("eagle", 3, "hello", 1, "world", 1));
+    Assert.assertEquals(wordCount, ImmutableMap.of("eagle", 3, "hello", 1));
   }
 
   private static class WordAndCount implements Serializable {