Initial commit for Ray streaming (#4268)

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

@@ -0,0 +1,75 @@
+package org.ray.streaming.demo;
+
+import com.google.common.collect.ImmutableMap;
+import org.ray.streaming.api.context.StreamingContext;
+import org.ray.streaming.api.function.impl.FlatMapFunction;
+import org.ray.streaming.api.function.impl.ReduceFunction;
+import org.ray.streaming.api.function.impl.SinkFunction;
+import org.ray.streaming.api.stream.StreamSource;
+import org.ray.streaming.util.ConfigKey;
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+
+public class WordCountTest implements Serializable {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(WordCountTest.class);
+
+  // TODO(zhenxuanpan): this test only works in single-process mode, because we put
+  //   results in this in-memory map.
+  static Map<String, Integer> wordCount = new ConcurrentHashMap<>();
+
+  @Test
+  public void testWordCount() {
+    StreamingContext streamingContext = StreamingContext.buildContext();
+    Map<String, Object> config = new HashMap<>();
+    config.put(ConfigKey.STREAMING_MAX_BATCH_COUNT, 1);
+    streamingContext.withConfig(config);
+    List<String> text = new ArrayList<>();
+    text.add("hello world eagle eagle eagle");
+    StreamSource<String> streamSource = StreamSource.buildSource(streamingContext, text);
+    streamSource
+        .flatMap((FlatMapFunction<String, WordAndCount>) (value, collector) -> {
+          String[] records = value.split(" ");
+          for (String record : records) {
+            collector.collect(new WordAndCount(record, 1));
+          }
+        })
+        .keyBy(pair -> pair.word)
+        .reduce((ReduceFunction<WordAndCount>) (oldValue, newValue) ->
+            new WordAndCount(oldValue.word, oldValue.count + newValue.count))
+        .sink((SinkFunction<WordAndCount>) result -> wordCount.put(result.word, result.count));
+
+    streamingContext.execute();
+
+    // Sleep until the count for every word is computed.
+    while (wordCount.size() < 3) {
+      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));
+  }
+
+  private static class WordAndCount implements Serializable {
+
+    public final String word;
+    public final Integer count;
+
+    public WordAndCount(String key, Integer count) {
+      this.word = key;
+      this.count = count;
+    }
+  }
+
+}

java/streaming/src/test/java/org/ray/streaming/plan/PlanBuilderTest.java

@@ -0,0 +1,87 @@
+package org.ray.streaming.plan;
+
+
+import com.google.common.collect.Lists;
+import org.ray.streaming.api.context.StreamingContext;
+import org.ray.streaming.api.partition.impl.KeyPartition;
+import org.ray.streaming.api.partition.impl.RoundRobinPartition;
+import org.ray.streaming.api.stream.DataStream;
+import org.ray.streaming.api.stream.StreamSink;
+import org.ray.streaming.api.stream.StreamSource;
+import java.util.List;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+public class PlanBuilderTest {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(PlanBuilderTest.class);
+
+  @Test
+  public void testDataSync() {
+    Plan plan = buildDataSyncPlan();
+    List<PlanVertex> planVertexList = plan.getPlanVertexList();
+    List<PlanEdge> planEdgeList = plan.getPlanEdgeList();
+
+    Assert.assertEquals(planVertexList.size(), 2);
+    Assert.assertEquals(planEdgeList.size(), 1);
+
+    PlanEdge planEdge = planEdgeList.get(0);
+    Assert.assertEquals(planEdge.getPartition().getClass(), RoundRobinPartition.class);
+
+    PlanVertex sinkVertex = planVertexList.get(1);
+    PlanVertex sourceVertex = planVertexList.get(0);
+    Assert.assertEquals(sinkVertex.getVertexType(), VertexType.SINK);
+    Assert.assertEquals(sourceVertex.getVertexType(), VertexType.SOURCE);
+
+  }
+
+  public Plan buildDataSyncPlan() {
+    StreamingContext streamingContext = StreamingContext.buildContext();
+    DataStream<String> dataStream = StreamSource.buildSource(streamingContext,
+        Lists.newArrayList("a", "b", "c"));
+    StreamSink streamSink = dataStream.sink(x -> LOGGER.info(x));
+    PlanBuilder planBuilder = new PlanBuilder(Lists.newArrayList(streamSink));
+
+    Plan plan = planBuilder.buildPlan();
+    return plan;
+  }
+
+  @Test
+  public void testKeyByPlan() {
+    Plan plan = buildKeyByPlan();
+    List<PlanVertex> planVertexList = plan.getPlanVertexList();
+    List<PlanEdge> planEdgeList = plan.getPlanEdgeList();
+
+    Assert.assertEquals(planVertexList.size(), 3);
+    Assert.assertEquals(planEdgeList.size(), 2);
+
+    PlanVertex source = planVertexList.get(0);
+    PlanVertex map = planVertexList.get(1);
+    PlanVertex sink = planVertexList.get(2);
+
+    Assert.assertEquals(source.getVertexType(), VertexType.SOURCE);
+    Assert.assertEquals(map.getVertexType(), VertexType.PROCESS);
+    Assert.assertEquals(sink.getVertexType(), VertexType.SINK);
+
+    PlanEdge keyBy2Sink = planEdgeList.get(0);
+    PlanEdge source2KeyBy = planEdgeList.get(1);
+
+    Assert.assertEquals(keyBy2Sink.getPartition().getClass(), KeyPartition.class);
+    Assert.assertEquals(source2KeyBy.getPartition().getClass(), RoundRobinPartition.class);
+  }
+
+  public Plan buildKeyByPlan() {
+    StreamingContext streamingContext = StreamingContext.buildContext();
+    DataStream<String> dataStream = StreamSource.buildSource(streamingContext,
+        Lists.newArrayList("1", "2", "3", "4"));
+    StreamSink streamSink = dataStream.keyBy(x -> x)
+        .sink(x -> LOGGER.info(x));
+    PlanBuilder planBuilder = new PlanBuilder(Lists.newArrayList(streamSink));
+
+    Plan plan = planBuilder.buildPlan();
+    return plan;
+  }
+
+}

java/streaming/src/test/java/org/ray/streaming/schedule/impl/TaskAssignImplTest.java

@@ -0,0 +1,58 @@
+package org.ray.streaming.schedule.impl;
+
+import org.ray.streaming.api.partition.impl.RoundRobinPartition;
+import org.ray.streaming.core.graph.ExecutionEdge;
+import org.ray.streaming.core.graph.ExecutionGraph;
+import org.ray.streaming.core.graph.ExecutionNode;
+import org.ray.streaming.core.graph.ExecutionNode.NodeType;
+import org.ray.streaming.core.runtime.StreamWorker;
+import org.ray.streaming.plan.Plan;
+import org.ray.streaming.plan.PlanBuilderTest;
+import org.ray.streaming.schedule.ITaskAssign;
+import java.util.ArrayList;
+import java.util.List;
+import org.ray.api.RayActor;
+import org.ray.runtime.RayActorImpl;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+public class TaskAssignImplTest {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(TaskAssignImplTest.class);
+
+  @Test
+  public void testTaskAssignImpl() {
+    PlanBuilderTest planBuilderTest = new PlanBuilderTest();
+    Plan plan = planBuilderTest.buildDataSyncPlan();
+
+    List<RayActor<StreamWorker>> workers = new ArrayList<>();
+    for(int i = 0; i < plan.getPlanVertexList().size(); i++) {
+      workers.add(new RayActorImpl<>());
+    }
+
+    ITaskAssign taskAssign = new TaskAssignImpl();
+    ExecutionGraph executionGraph = taskAssign.assign(plan, workers);
+
+    List<ExecutionNode> executionNodeList = executionGraph.getExecutionNodeList();
+
+    Assert.assertEquals(executionNodeList.size(), 2);
+    ExecutionNode sourceNode = executionNodeList.get(0);
+    Assert.assertEquals(sourceNode.getNodeType(), NodeType.SOURCE);
+    Assert.assertEquals(sourceNode.getExecutionTaskList().size(), 1);
+    Assert.assertEquals(sourceNode.getExecutionEdgeList().size(), 1);
+
+    List<ExecutionEdge> sourceExecutionEdges = sourceNode.getExecutionEdgeList();
+
+    Assert.assertEquals(sourceExecutionEdges.size(), 1);
+    ExecutionEdge source2Sink = sourceExecutionEdges.get(0);
+
+    Assert.assertEquals(source2Sink.getPartition().getClass(), RoundRobinPartition.class);
+
+    ExecutionNode sinkNode = executionNodeList.get(1);
+    Assert.assertEquals(sinkNode.getNodeType(), NodeType.SINK);
+    Assert.assertEquals(sinkNode.getExecutionTaskList().size(), 1);
+    Assert.assertEquals(sinkNode.getExecutionEdgeList().size(), 0);
+  }
+}

java/streaming/src/test/resources/log4j.properties

@@ -0,0 +1,6 @@
+log4j.rootLogger=INFO, stdout
+# Direct log messages to stdout
+log4j.appender.stdout=org.apache.log4j.ConsoleAppender
+log4j.appender.stdout.Target=System.out
+log4j.appender.stdout.layout=org.apache.log4j.PatternLayout
+log4j.appender.stdout.layout.ConversionPattern=%d{yyyy-MM-dd HH:mm:ss} %-5p %c{1}:%L - %m%n

java/streaming/src/test/resources/ray.conf

@@ -0,0 +1,3 @@
+ray {
+  run-mode = SINGLE_PROCESS
+}