[Streaming] Union api (#8612)

streaming/java/streaming-runtime/src/main/java/io/ray/streaming/runtime/python/GraphPbBuilder.java

@@ -4,7 +4,9 @@ import com.google.protobuf.ByteString;
 import io.ray.runtime.actor.NativeRayActor;
 import io.ray.streaming.api.function.Function;
 import io.ray.streaming.api.partition.Partition;
+import io.ray.streaming.operator.Operator;
 import io.ray.streaming.python.PythonFunction;
+import io.ray.streaming.python.PythonOperator;
 import io.ray.streaming.python.PythonPartition;
 import io.ray.streaming.runtime.core.graph.ExecutionEdge;
 import io.ray.streaming.runtime.core.graph.ExecutionGraph;
@@ -34,8 +36,8 @@ public class GraphPbBuilder {
       nodeBuilder.setNodeType(
           Streaming.NodeType.valueOf(node.getNodeType().name()));
       nodeBuilder.setLanguage(Streaming.Language.valueOf(node.getLanguage().name()));
-      byte[] functionBytes = serializeFunction(node.getStreamOperator().getFunction());
-      nodeBuilder.setFunction(ByteString.copyFrom(functionBytes));
+      byte[] operatorBytes = serializeOperator(node.getStreamOperator());
+      nodeBuilder.setOperator(ByteString.copyFrom(operatorBytes));
 
       // build tasks
       for (ExecutionTask task : node.getExecutionTasks()) {
@@ -72,6 +74,19 @@ public class GraphPbBuilder {
     return edgeBuilder.build();
   }
 
+  private byte[] serializeOperator(Operator operator) {
+    if (operator instanceof PythonOperator) {
+      PythonOperator pythonOperator = (PythonOperator) operator;
+      return serializer.serialize(Arrays.asList(
+          serializeFunction(pythonOperator.getFunction()),
+          pythonOperator.getModuleName(),
+          pythonOperator.getClassName()
+      ));
+    } else {
+      return new byte[0];
+    }
+  }
+
   private byte[] serializeFunction(Function function) {
     if (function instanceof PythonFunction) {
       PythonFunction pyFunc = (PythonFunction) function;

streaming/java/streaming-runtime/src/main/java/io/ray/streaming/runtime/python/PythonGateway.java

@@ -3,8 +3,11 @@ package io.ray.streaming.runtime.python;
 import com.google.common.base.Preconditions;
 import com.google.common.primitives.Primitives;
 import io.ray.streaming.api.context.StreamingContext;
+import io.ray.streaming.api.stream.DataStream;
+import io.ray.streaming.api.stream.Stream;
 import io.ray.streaming.python.PythonFunction;
 import io.ray.streaming.python.PythonPartition;
+import io.ray.streaming.python.stream.PythonDataStream;
 import io.ray.streaming.python.stream.PythonStreamSource;
 import io.ray.streaming.runtime.serialization.MsgPackSerializer;
 import io.ray.streaming.runtime.util.ReflectionUtils;
@@ -99,6 +102,26 @@ public class PythonGateway {
     return serializer.serialize(getReferenceId(partition));
   }
 
+  public byte[] union(byte[] paramsBytes) {
+    List<Object> streams = (List<Object>) serializer.deserialize(paramsBytes);
+    streams = processParameters(streams);
+    LOG.info("Call union with streams {}", streams);
+    Preconditions.checkArgument(streams.size() >= 2,
+        "Union needs at least two streams");
+    Stream unionStream;
+    Stream stream1 = (Stream) streams.get(0);
+    List otherStreams = streams.subList(1, streams.size());
+    if (stream1 instanceof DataStream) {
+      DataStream dataStream = (DataStream) stream1;
+      unionStream = dataStream.union(otherStreams);
+    } else {
+      Preconditions.checkArgument(stream1 instanceof PythonDataStream);
+      PythonDataStream pythonDataStream = (PythonDataStream) stream1;
+      unionStream = pythonDataStream.union(otherStreams);
+    }
+    return serialize(unionStream);
+  }
+
   public byte[] callFunction(byte[] paramsBytes) {
     try {
       List<Object> params = (List<Object>) serializer.deserialize(paramsBytes);
@@ -111,12 +134,7 @@ public class PythonGateway {
           .map(Object::getClass).toArray(Class[]::new);
       Method method = findMethod(clz, funcName, paramsTypes);
       Object result = method.invoke(null, params.subList(2, params.size()).toArray());
-      if (returnReference(result)) {
-        referenceMap.put(getReferenceId(result), result);
-        return serializer.serialize(getReferenceId(result));
-      } else {
-        return serializer.serialize(result);
-      }
+      return serialize(result);
     } catch (Exception e) {
       throw new RuntimeException(e);
     }
@@ -134,12 +152,7 @@ public class PythonGateway {
           .map(Object::getClass).toArray(Class[]::new);
       Method method = findMethod(clz, methodName, paramsTypes);
       Object result = method.invoke(obj, params.subList(2, params.size()).toArray());
-      if (returnReference(result)) {
-        referenceMap.put(getReferenceId(result), result);
-        return serializer.serialize(getReferenceId(result));
-      } else {
-        return serializer.serialize(result);
-      }
+      return serialize(result);
     } catch (Exception e) {
       throw new RuntimeException(e);
     }
@@ -179,6 +192,15 @@ public class PythonGateway {
     return any.get();
   }
 
+  private byte[] serialize(Object value) {
+    if (returnReference(value)) {
+      referenceMap.put(getReferenceId(value), value);
+      return serializer.serialize(getReferenceId(value));
+    } else {
+      return serializer.serialize(value);
+    }
+  }
+
   private static boolean returnReference(Object value) {
     if (isBasic(value)) {
       return false;

streaming/java/streaming-runtime/src/test/java/io/ray/streaming/runtime/demo/HybridStreamTest.java

@@ -4,16 +4,22 @@ import io.ray.api.Ray;
 import io.ray.streaming.api.context.StreamingContext;
 import io.ray.streaming.api.function.impl.FilterFunction;
 import io.ray.streaming.api.function.impl.MapFunction;
+import io.ray.streaming.api.function.impl.SinkFunction;
 import io.ray.streaming.api.stream.DataStreamSource;
 import io.ray.streaming.runtime.BaseUnitTest;
+import java.io.IOException;
 import java.io.Serializable;
+import java.nio.file.Files;
+import java.nio.file.Paths;
+import java.nio.file.StandardOpenOption;
 import java.util.Arrays;
 import java.util.concurrent.TimeUnit;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.testng.Assert;
 import org.testng.annotations.Test;
 
-public class HybridStreamTest extends BaseUnitTest implements Serializable {
+public class HybridStreamTest {
   private static final Logger LOG = LoggerFactory.getLogger(HybridStreamTest.class);
 
   public static class Mapper1 implements MapFunction<Object, Object> {
@@ -34,9 +40,12 @@ public class HybridStreamTest extends BaseUnitTest implements Serializable {
     }
   }
 
-  @Test
-  public void testHybridDataStream() throws InterruptedException {
+  @Test(timeOut = 60000)
+  public void testHybridDataStream() throws Exception {
     Ray.shutdown();
+    String sinkFileName = "/tmp/testHybridDataStream.txt";
+    Files.deleteIfExists(Paths.get(sinkFileName));
+
     StreamingContext context = StreamingContext.buildContext();
     DataStreamSource<String> streamSource =
         DataStreamSource.fromCollection(context, Arrays.asList("a", "b", "c"));
@@ -46,9 +55,38 @@ public class HybridStreamTest extends BaseUnitTest implements Serializable {
         .map("ray.streaming.tests.test_hybrid_stream", "map_func1")
         .filter("ray.streaming.tests.test_hybrid_stream", "filter_func1")
         .asJavaStream()
-        .sink(x -> System.out.println("HybridStreamTest: " + x));
+        .sink((SinkFunction<Object>) value -> {
+          LOG.info("HybridStreamTest: {}", value);
+          try {
+            if (!Files.exists(Paths.get(sinkFileName))) {
+              Files.createFile(Paths.get(sinkFileName));
+            }
+            Files.write(Paths.get(sinkFileName), value.toString().getBytes(),
+                StandardOpenOption.APPEND);
+          } catch (IOException e) {
+            throw new RuntimeException(e);
+          }
+        });
     context.execute("HybridStreamTestJob");
+    int sleptTime = 0;
     TimeUnit.SECONDS.sleep(3);
+    while (true) {
+      if (Files.exists(Paths.get(sinkFileName))) {
+        TimeUnit.SECONDS.sleep(3);
+        String text = String.join(", ", Files.readAllLines(Paths.get(sinkFileName)));
+        Assert.assertTrue(text.contains("a"));
+        Assert.assertFalse(text.contains("b"));
+        Assert.assertTrue(text.contains("c"));
+        LOG.info("Execution succeed");
+        break;
+      }
+      sleptTime += 1;
+      if (sleptTime >= 60) {
+        throw new RuntimeException("Execution not finished");
+      }
+      LOG.info("Wait finish...");
+      TimeUnit.SECONDS.sleep(1);
+    }
     context.stop();
     LOG.info("HybridStreamTest succeed");
   }

streaming/java/streaming-runtime/src/test/java/io/ray/streaming/runtime/demo/UnionStreamTest.java

@@ -0,0 +1,71 @@
+package io.ray.streaming.runtime.demo;
+
+import io.ray.api.Ray;
+import io.ray.streaming.api.context.StreamingContext;
+import io.ray.streaming.api.function.impl.SinkFunction;
+import io.ray.streaming.api.stream.DataStreamSource;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Paths;
+import java.nio.file.StandardOpenOption;
+import java.util.Arrays;
+import java.util.concurrent.TimeUnit;
+import org.apache.commons.lang3.StringUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+public class UnionStreamTest {
+  private static final Logger LOG = LoggerFactory.getLogger( UnionStreamTest.class );
+
+  @Test(timeOut = 60000)
+  public void testUnionStream() throws Exception {
+    Ray.shutdown();
+    String sinkFileName = "/tmp/testUnionStream.txt";
+    Files.deleteIfExists(Paths.get(sinkFileName));
+
+    StreamingContext context = StreamingContext.buildContext();
+    DataStreamSource<Integer> streamSource1 =
+        DataStreamSource.fromCollection(context, Arrays.asList(1, 1));
+    DataStreamSource<Integer> streamSource2 =
+        DataStreamSource.fromCollection(context, Arrays.asList(1, 1));
+    DataStreamSource<Integer> streamSource3 =
+        DataStreamSource.fromCollection(context, Arrays.asList(1, 1));
+    streamSource1
+        .union(streamSource2, streamSource3)
+        .sink((SinkFunction<Integer>) value -> {
+          LOG.info("UnionStreamTest: {}", value);
+          try {
+            if (!Files.exists(Paths.get(sinkFileName))) {
+              Files.createFile(Paths.get(sinkFileName));
+            }
+            Files.write(Paths.get(sinkFileName), value.toString().getBytes(),
+                StandardOpenOption.APPEND);
+          } catch (IOException e) {
+            throw new RuntimeException(e);
+          }
+        });
+    context.execute("UnionStreamTest");
+    int sleptTime = 0;
+    TimeUnit.SECONDS.sleep(3);
+    while (true) {
+      if (Files.exists(Paths.get(sinkFileName))) {
+        TimeUnit.SECONDS.sleep(3);
+        String text = String.join(", ", Files.readAllLines(Paths.get(sinkFileName)));
+        Assert.assertEquals(text, StringUtils.repeat("1", 6));
+        LOG.info("Execution succeed");
+        break;
+      }
+      sleptTime += 1;
+      if (sleptTime >= 60) {
+        throw new RuntimeException("Execution not finished");
+      }
+      LOG.info("Wait finish...");
+      TimeUnit.SECONDS.sleep(1);
+    }
+    context.stop();
+    LOG.info("HybridStreamTest succeed");
+  }
+
+}