[Streaming] Streaming Cross-Lang API (#7464)

streaming/python/datastream.py

@@ -1,4 +1,4 @@
-from abc import ABC
+from abc import ABC, abstractmethod
 
 from ray.streaming import function
 from ray.streaming import partition
@@ -19,7 +19,6 @@ class Stream(ABC):
             self.streaming_context = input_stream.streaming_context
         else:
             self.streaming_context = streaming_context
-        self.parallelism = 1
 
     def get_streaming_context(self):
         return self.streaming_context
@@ -29,7 +28,8 @@ class Stream(ABC):
         Returns:
             the parallelism of this transformation
         """
-        return self.parallelism
+        return self._gateway_client(). \
+            call_method(self._j_stream, "getParallelism")
 
     def set_parallelism(self, parallelism: int):
         """Sets the parallelism of this transformation
@@ -40,7 +40,6 @@ class Stream(ABC):
         Returns:
             self
         """
-        self.parallelism = parallelism
         self._gateway_client(). \
             call_method(self._j_stream, "setParallelism", parallelism)
         return self
@@ -60,6 +59,10 @@ class Stream(ABC):
         return self._gateway_client(). \
             call_method(self._j_stream, "getId")
 
+    @abstractmethod
+    def get_language(self):
+        pass
+
     def _gateway_client(self):
         return self.get_streaming_context()._gateway_client
 
@@ -75,6 +78,9 @@ class DataStream(Stream):
         super().__init__(
             input_stream, j_stream, streaming_context=streaming_context)
 
+    def get_language(self):
+        return function.Language.PYTHON
+
     def map(self, func):
         """
         Applies a Map transformation on a :class:`DataStream`.
@@ -158,6 +164,7 @@ class DataStream(Stream):
         Returns:
              A KeyDataStream
         """
+        self._check_partition_call()
         if not isinstance(func, function.KeyFunction):
             func = function.SimpleKeyFunction(func)
         j_func = self._gateway_client().create_py_func(
@@ -175,6 +182,7 @@ class DataStream(Stream):
         Returns:
             The DataStream with broadcast partitioning set.
         """
+        self._check_partition_call()
         self._gateway_client().call_method(self._j_stream, "broadcast")
         return self
 
@@ -191,6 +199,7 @@ class DataStream(Stream):
         Returns:
             The DataStream with specified partitioning set.
         """
+        self._check_partition_call()
         if not isinstance(partition_func, partition.Partition):
             partition_func = partition.SimplePartition(partition_func)
         j_partition = self._gateway_client().create_py_func(
@@ -199,6 +208,16 @@ class DataStream(Stream):
             call_method(self._j_stream, "partitionBy", j_partition)
         return self
 
+    def _check_partition_call(self):
+        """
+        If parent stream is a java stream, we can't call partition related
+        methods in the python stream
+        """
+        if self.input_stream is not None and \
+                self.input_stream.get_language() == function.Language.JAVA:
+            raise Exception("Partition related methods can't be called on a "
+                            "python stream if parent stream is a java stream.")
+
     def sink(self, func):
         """
         Create a StreamSink with the given sink.
@@ -217,8 +236,97 @@ class DataStream(Stream):
             call_method(self._j_stream, "sink", j_func)
         return StreamSink(self, j_stream, func)
 
+    def as_java_stream(self):
+        """
+        Convert this stream as a java JavaDataStream.
+        The converted stream and this stream are the same logical stream,
+        which has same stream id. Changes in converted stream will be reflected
+        in this stream and vice versa.
+        """
+        j_stream = self._gateway_client(). \
+            call_method(self._j_stream, "asJavaStream")
+        return JavaDataStream(self, j_stream)
 
-class KeyDataStream(Stream):
+
+class JavaDataStream(Stream):
+    """
+    Represents a stream of data which applies a transformation executed by
+    java. It's also a wrapper of java
+    `org.ray.streaming.api.stream.DataStream`
+    """
+
+    def __init__(self, input_stream, j_stream, streaming_context=None):
+        super().__init__(
+            input_stream, j_stream, streaming_context=streaming_context)
+
+    def get_language(self):
+        return function.Language.JAVA
+
+    def map(self, java_func_class):
+        """See org.ray.streaming.api.stream.DataStream.map"""
+        return JavaDataStream(self, self._unary_call("map", java_func_class))
+
+    def flat_map(self, java_func_class):
+        """See org.ray.streaming.api.stream.DataStream.flatMap"""
+        return JavaDataStream(self, self._unary_call("flatMap",
+                                                     java_func_class))
+
+    def filter(self, java_func_class):
+        """See org.ray.streaming.api.stream.DataStream.filter"""
+        return JavaDataStream(self, self._unary_call("filter",
+                                                     java_func_class))
+
+    def key_by(self, java_func_class):
+        """See org.ray.streaming.api.stream.DataStream.keyBy"""
+        self._check_partition_call()
+        return JavaKeyDataStream(self,
+                                 self._unary_call("keyBy", java_func_class))
+
+    def broadcast(self, java_func_class):
+        """See org.ray.streaming.api.stream.DataStream.broadcast"""
+        self._check_partition_call()
+        return JavaDataStream(self,
+                              self._unary_call("broadcast", java_func_class))
+
+    def partition_by(self, java_func_class):
+        """See org.ray.streaming.api.stream.DataStream.partitionBy"""
+        self._check_partition_call()
+        return JavaDataStream(self,
+                              self._unary_call("partitionBy", java_func_class))
+
+    def sink(self, java_func_class):
+        """See org.ray.streaming.api.stream.DataStream.sink"""
+        return JavaStreamSink(self, self._unary_call("sink", java_func_class))
+
+    def as_python_stream(self):
+        """
+        Convert this stream as a python DataStream.
+        The converted stream and this stream are the same logical stream,
+        which has same stream id. Changes in converted stream will be reflected
+        in this stream and vice versa.
+        """
+        j_stream = self._gateway_client(). \
+            call_method(self._j_stream, "asPythonStream")
+        return DataStream(self, j_stream)
+
+    def _check_partition_call(self):
+        """
+        If parent stream is a python stream, we can't call partition related
+        methods in the java stream
+        """
+        if self.input_stream is not None and \
+                self.input_stream.get_language() == function.Language.PYTHON:
+            raise Exception("Partition related methods can't be called on a"
+                            "java stream if parent stream is a python stream.")
+
+    def _unary_call(self, func_name, java_func_class):
+        j_func = self._gateway_client().new_instance(java_func_class)
+        j_stream = self._gateway_client(). \
+            call_method(self._j_stream, func_name, j_func)
+        return j_stream
+
+
+class KeyDataStream(DataStream):
     """Represents a DataStream returned by a key-by operation.
      Wrapper of java io.ray.streaming.python.stream.PythonKeyDataStream
     """
@@ -251,6 +359,43 @@ class KeyDataStream(Stream):
             call_method(self._j_stream, "reduce", j_func)
         return DataStream(self, j_stream)
 
+    def as_java_stream(self):
+        """
+        Convert this stream as a java KeyDataStream.
+        The converted stream and this stream are the same logical stream,
+        which has same stream id. Changes in converted stream will be reflected
+        in this stream and vice versa.
+        """
+        j_stream = self._gateway_client(). \
+            call_method(self._j_stream, "asJavaStream")
+        return JavaKeyDataStream(self, j_stream)
+
+
+class JavaKeyDataStream(JavaDataStream):
+    """
+    Represents a DataStream returned by a key-by operation in java.
+     Wrapper of org.ray.streaming.api.stream.KeyDataStream
+    """
+
+    def __init__(self, input_stream, j_stream):
+        super().__init__(input_stream, j_stream)
+
+    def reduce(self, java_func_class):
+        """See org.ray.streaming.api.stream.KeyDataStream.reduce"""
+        return JavaDataStream(self,
+                              super()._unary_call("reduce", java_func_class))
+
+    def as_python_stream(self):
+        """
+        Convert this stream as a python KeyDataStream.
+        The converted stream and this stream are the same logical stream,
+        which has same stream id. Changes in converted stream will be reflected
+        in this stream and vice versa.
+        """
+        j_stream = self._gateway_client(). \
+            call_method(self._j_stream, "asPythonStream")
+        return KeyDataStream(self, j_stream)
+
 
 class StreamSource(DataStream):
     """Represents a source of the DataStream.
@@ -261,9 +406,12 @@ class StreamSource(DataStream):
         super().__init__(None, j_stream, streaming_context=streaming_context)
         self.source_func = source_func
 
+    def get_language(self):
+        return function.Language.PYTHON
+
     @staticmethod
     def build_source(streaming_context, func):
-        """Build a StreamSource source from a collection.
+        """Build a StreamSource source from a source function.
         Args:
             streaming_context: Stream context
             func: A instance of `SourceFunction`
@@ -275,6 +423,34 @@ class StreamSource(DataStream):
         return StreamSource(j_stream, streaming_context, func)
 
 
+class JavaStreamSource(JavaDataStream):
+    """Represents a source of the java DataStream.
+     Wrapper of java org.ray.streaming.api.stream.DataStreamSource
+    """
+
+    def __init__(self, j_stream, streaming_context):
+        super().__init__(None, j_stream, streaming_context=streaming_context)
+
+    def get_language(self):
+        return function.Language.JAVA
+
+    @staticmethod
+    def build_source(streaming_context, java_source_func_class):
+        """Build a java StreamSource source from a java source function.
+        Args:
+            streaming_context: Stream context
+            java_source_func_class: qualified class name of java SourceFunction
+        Returns:
+            A java StreamSource
+        """
+        j_func = streaming_context._gateway_client() \
+            .new_instance(java_source_func_class)
+        j_stream = streaming_context._gateway_client() \
+            .call_function("org.ray.streaming.api.stream.DataStreamSource"
+                           "fromSource", streaming_context._j_ctx, j_func)
+        return JavaStreamSource(j_stream, streaming_context)
+
+
 class StreamSink(Stream):
     """Represents a sink of the DataStream.
      Wrapper of java io.ray.streaming.python.stream.PythonStreamSink
@@ -282,3 +458,18 @@ class StreamSink(Stream):
 
     def __init__(self, input_stream, j_stream, func):
         super().__init__(input_stream, j_stream)
+
+    def get_language(self):
+        return function.Language.PYTHON
+
+
+class JavaStreamSink(Stream):
+    """Represents a sink of the java DataStream.
+     Wrapper of java org.ray.streaming.api.stream.StreamSink
+    """
+
+    def __init__(self, input_stream, j_stream):
+        super().__init__(input_stream, j_stream)
+
+    def get_language(self):
+        return function.Language.JAVA