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

streaming/python/runtime/serialization.py

@@ -0,0 +1,57 @@
+from abc import ABC, abstractmethod
+import pickle
+import msgpack
+from ray.streaming import message
+
+_RECORD_TYPE_ID = 0
+_KEY_RECORD_TYPE_ID = 1
+_CROSS_LANG_TYPE_ID = b"0"
+_JAVA_TYPE_ID = b"1"
+_PYTHON_TYPE_ID = b"2"
+
+
+class Serializer(ABC):
+    @abstractmethod
+    def serialize(self, obj):
+        pass
+
+    @abstractmethod
+    def deserialize(self, serialized_bytes):
+        pass
+
+
+class PythonSerializer(Serializer):
+    def serialize(self, obj):
+        return pickle.dumps(obj)
+
+    def deserialize(self, serialized_bytes):
+        return pickle.loads(serialized_bytes)
+
+
+class CrossLangSerializer(Serializer):
+    """Serialize stream element between java/python"""
+
+    def serialize(self, obj):
+        if type(obj) is message.Record:
+            fields = [_RECORD_TYPE_ID, obj.stream, obj.value]
+        elif type(obj) is message.KeyRecord:
+            fields = [_KEY_RECORD_TYPE_ID, obj.stream, obj.key, obj.value]
+        else:
+            raise Exception("Unsupported value {}".format(obj))
+        return msgpack.packb(fields, use_bin_type=True)
+
+    def deserialize(self, data):
+        fields = msgpack.unpackb(data, raw=False)
+        if fields[0] == _RECORD_TYPE_ID:
+            stream, value = fields[1:]
+            record = message.Record(value)
+            record.stream = stream
+            return record
+        elif fields[0] == _KEY_RECORD_TYPE_ID:
+            stream, key, value = fields[1:]
+            key_record = message.KeyRecord(key, value)
+            key_record.stream = stream
+            return key_record
+        else:
+            raise Exception("Unsupported type id {}, type {}".format(
+                fields[0], type(fields[0])))