[Streaming] Union api (#8612)

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