[Streaming] Union api (#8612)

streaming/python/context.py

@@ -33,7 +33,7 @@ class StreamingContext:
                 self
             """
             if key is not None:
-                assert value
+                assert value is not None
                 self._options[key] = str(value)
             if conf is not None:
                 for k, v in conf.items():

streaming/python/datastream.py

@@ -71,7 +71,6 @@ class Stream(ABC):
             self
         """
         if key is not None:
-            assert value
             assert type(key) is str
             assert type(value) is str
             self._gateway_client(). \
@@ -183,6 +182,21 @@ class DataStream(Stream):
             call_method(self._j_stream, "filter", j_func)
         return DataStream(self, j_stream)
 
+    def union(self, *streams):
+        """Apply union transformations to this stream by merging data stream
+         outputs of the same type with each other.
+
+        Args:
+            *streams: The DataStreams to union output with.
+
+        Returns:
+            A new UnionStream.
+        """
+        assert len(streams) >= 1, "Need at least one stream to union with"
+        j_streams = [s._j_stream for s in streams]
+        j_stream = self._gateway_client().union(self._j_stream, *j_streams)
+        return UnionStream(self, j_stream)
+
     def key_by(self, func):
         """
         Creates a new :class:`KeyDataStream` that uses the provided key to
@@ -308,6 +322,13 @@ class JavaDataStream(Stream):
         return JavaDataStream(self, self._unary_call("filter",
                                                      java_func_class))
 
+    def union(self, *streams):
+        """See io.ray.streaming.api.stream.DataStream.union"""
+        assert len(streams) >= 1, "Need at least one stream to union with"
+        j_streams = [s._j_stream for s in streams]
+        j_stream = self._gateway_client().union(self._j_stream, *j_streams)
+        return JavaUnionStream(self, j_stream)
+
     def key_by(self, java_func_class):
         """See io.ray.streaming.api.stream.DataStream.keyBy"""
         self._check_partition_call()
@@ -429,6 +450,30 @@ class JavaKeyDataStream(JavaDataStream):
         return KeyDataStream(self, j_stream)
 
 
+class UnionStream(DataStream):
+    """Represents a union stream.
+     Wrapper of java io.ray.streaming.python.stream.PythonUnionStream
+    """
+
+    def __init__(self, input_stream, j_stream):
+        super().__init__(input_stream, j_stream)
+
+    def get_language(self):
+        return function.Language.PYTHON
+
+
+class JavaUnionStream(JavaDataStream):
+    """Represents a java union stream.
+     Wrapper of java io.ray.streaming.api.stream.UnionStream
+    """
+
+    def __init__(self, input_stream, j_stream):
+        super().__init__(input_stream, j_stream)
+
+    def get_language(self):
+        return function.Language.JAVA
+
+
 class StreamSource(DataStream):
     """Represents a source of the DataStream.
      Wrapper of java io.ray.streaming.python.stream.PythonStreamSource

streaming/python/function.py

@@ -23,6 +23,16 @@ class Function(ABC):
         pass
 
 
+class EmptyFunction(Function):
+    """Default function which does nothing"""
+
+    def open(self, runtime_context):
+        pass
+
+    def close(self):
+        pass
+
+
 class SourceContext(ABC):
     """
     Interface that source functions use to emit elements, and possibly
@@ -216,7 +226,7 @@ class SimpleFlatMapFunction(FlatMapFunction):
         self.func = func
         self.process_func = None
         sig = inspect.signature(func)
-        assert len(sig.parameters) <= 2,\
+        assert len(sig.parameters) <= 2, \
             "func should receive value [, collector] as arguments"
         if len(sig.parameters) == 2:
 
@@ -292,7 +302,7 @@ def load_function(descriptor_func_bytes: bytes):
         a streaming function
     """
     assert len(descriptor_func_bytes) > 0
-    function_bytes, module_name, function_name, function_interface\
+    function_bytes, module_name, function_name, function_interface \
         = gateway_client.deserialize(descriptor_func_bytes)
     if function_bytes:
         return deserialize(function_bytes)

streaming/python/operator.py

@@ -1,8 +1,11 @@
-from abc import ABC, abstractmethod
 import enum
+import importlib
+from abc import ABC, abstractmethod
+
 from ray import streaming
 from ray.streaming import function
 from ray.streaming import message
+from ray.streaming.runtime import gateway_client
 
 
 class OperatorType(enum.Enum):
@@ -214,6 +217,16 @@ class SinkOperator(StreamOperator, OneInputOperator):
         self.func.sink(record.value)
 
 
+class UnionOperator(StreamOperator, OneInputOperator):
+    """Operator for union operation"""
+
+    def __init__(self):
+        super().__init__(function.EmptyFunction())
+
+    def process_element(self, record):
+        self.collect(record)
+
+
 _function_to_operator = {
     function.SourceFunction: SourceOperator,
     function.MapFunction: MapOperator,
@@ -225,7 +238,36 @@ _function_to_operator = {
 }
 
 
-def create_operator(func: function.Function):
+def load_operator(descriptor_operator_bytes: bytes):
+    """
+    Deserialize `descriptor_operator_bytes` to get operator info, then
+    create streaming operator.
+    Note that this function must be kept in sync with
+     `io.ray.streaming.runtime.python.GraphPbBuilder.serializeOperator`
+
+    Args:
+        descriptor_operator_bytes: serialized operator info
+
+    Returns:
+        a streaming operator
+    """
+    assert len(descriptor_operator_bytes) > 0
+    function_desc_bytes, module_name, class_name \
+        = gateway_client.deserialize(descriptor_operator_bytes)
+    if function_desc_bytes:
+        return create_operator_with_func(
+            function.load_function(function_desc_bytes))
+    else:
+        assert module_name
+        assert class_name
+        mod = importlib.import_module(module_name)
+        cls = getattr(mod, class_name)
+        assert issubclass(cls, Operator)
+        print("cls", cls)
+        return cls()
+
+
+def create_operator_with_func(func: function.Function):
     """Create an operator according to a :class:`function.Function`
 
     Args:

streaming/python/runtime/gateway_client.py

@@ -30,7 +30,7 @@ class GatewayClient:
 
     def create_py_stream_source(self, serialized_func):
         assert isinstance(serialized_func, bytes)
-        call = self._python_gateway_actor.createPythonStreamSource\
+        call = self._python_gateway_actor.createPythonStreamSource \
             .remote(serialized_func)
         return deserialize(ray.get(call))
 
@@ -41,10 +41,16 @@ class GatewayClient:
 
     def create_py_partition(self, serialized_partition):
         assert isinstance(serialized_partition, bytes)
-        call = self._python_gateway_actor.createPyPartition\
+        call = self._python_gateway_actor.createPyPartition \
             .remote(serialized_partition)
         return deserialize(ray.get(call))
 
+    def union(self, *streams):
+        serialized_streams = serialize(streams)
+        call = self._python_gateway_actor.union \
+            .remote(serialized_streams)
+        return deserialize(ray.get(call))
+
     def call_function(self, java_class, java_function, *args):
         java_params = serialize([java_class, java_function] + list(args))
         call = self._python_gateway_actor.callFunction.remote(java_params)

streaming/python/runtime/graph.py

@@ -5,7 +5,6 @@ import ray.streaming.generated.remote_call_pb2 as remote_call_pb
 import ray.streaming.generated.streaming_pb2 as streaming_pb
 import ray.streaming.operator as operator
 import ray.streaming.partition as partition
-from ray.streaming import function
 from ray.streaming.generated.streaming_pb2 import Language
 
 
@@ -32,9 +31,8 @@ class ExecutionNode:
             node_pb.node_type)]
         self.parallelism = node_pb.parallelism
         if node_pb.language == Language.PYTHON:
-            func_bytes = node_pb.function  # python function descriptor
-            func = function.load_function(func_bytes)
-            self.stream_operator = operator.create_operator(func)
+            operator_bytes = node_pb.operator  # python operator descriptor
+            self.stream_operator = operator.load_operator(operator_bytes)
         self.execution_tasks = [
             ExecutionTask(task) for task in node_pb.execution_tasks
         ]

streaming/python/tests/test_operator.py

@@ -1,8 +1,37 @@
-from ray.streaming import operator
 from ray.streaming import function
+from ray.streaming import operator
+from ray.streaming.operator import OperatorType
+from ray.streaming.runtime import gateway_client
 
 
-def test_create_operator():
+def test_create_operator_with_func():
     map_func = function.SimpleMapFunction(lambda x: x)
-    map_operator = operator.create_operator(map_func)
+    map_operator = operator.create_operator_with_func(map_func)
     assert type(map_operator) is operator.MapOperator
+
+
+class MapFunc(function.MapFunction):
+    def map(self, value):
+        return str(value)
+
+
+class EmptyOperator(operator.StreamOperator):
+    def __init__(self):
+        super().__init__(function.EmptyFunction())
+
+    def operator_type(self) -> OperatorType:
+        return OperatorType.ONE_INPUT
+
+
+def test_load_operator():
+    # function_bytes, module_name, class_name,
+    descriptor_func_bytes = gateway_client.serialize(
+        [None, __name__, MapFunc.__name__, "MapFunction"])
+    descriptor_op_bytes = gateway_client.serialize(
+        [descriptor_func_bytes, "", ""])
+    map_operator = operator.load_operator(descriptor_op_bytes)
+    assert type(map_operator) is operator.MapOperator
+    descriptor_op_bytes = gateway_client.serialize(
+        [None, __name__, EmptyOperator.__name__])
+    test_operator = operator.load_operator(descriptor_op_bytes)
+    assert isinstance(test_operator, EmptyOperator)

streaming/python/tests/test_union_stream.py

@@ -0,0 +1,47 @@
+import os
+
+import ray
+from ray.streaming import StreamingContext
+
+
+def test_union_stream():
+    ray.init(load_code_from_local=True, include_java=True)
+    ctx = StreamingContext.Builder() \
+        .option("streaming.metrics.reporters", "") \
+        .build()
+    sink_file = "/tmp/test_union_stream.txt"
+    if os.path.exists(sink_file):
+        os.remove(sink_file)
+
+    def sink_func(x):
+        with open(sink_file, "a") as f:
+            print("sink_func", x)
+            f.write(str(x))
+
+    stream1 = ctx.from_values(1, 2)
+    stream2 = ctx.from_values(3, 4)
+    stream3 = ctx.from_values(5, 6)
+    stream1.union(stream2, stream3).sink(sink_func)
+    ctx.submit("test_union_stream")
+    import time
+    slept_time = 0
+    while True:
+        if os.path.exists(sink_file):
+            time.sleep(3)
+            with open(sink_file, "r") as f:
+                result = f.read()
+                print("sink result", result)
+                assert set(result) == {"1", "2", "3", "4", "5", "6"}
+            print("Execution succeed")
+            break
+        if slept_time >= 60:
+            raise Exception("Execution not finished")
+        slept_time = slept_time + 1
+        print("Wait finish...")
+        time.sleep(1)
+
+    ray.shutdown()
+
+
+if __name__ == "__main__":
+    test_union_stream()