Add FuseChunkData and FuseChunk to represent fused chunk (mars-project#412)

* Add FuseChunkData and FuseChunk to represent fused chunk, refactor serialization for all kinds of chunks

Author: Qin Xuye

mars/core.py

@@ -25,7 +25,7 @@
 from .utils import tokenize, AttributeDict, on_serialize_shape, \
     on_deserialize_shape, on_serialize_nsplits, is_eager_mode
 from .serialize import HasKey, ValueType, ProviderType, Serializable, AttributeAsDict, \
-    TupleField, DictField, KeyField, BoolField, StringField
+    TupleField, ListField, DictField, KeyField, BoolField, StringField, OneOfField
 from .tiles import Tileable, handler
 from .graph import DAG
 
@@ -281,6 +281,73 @@ class Chunk(Entity):
     _allow_data_type_ = (ChunkData,)
 
 
+def _on_serialize_composed(composed):
+    return [FuseChunkData.ChunkRef(c.data if isinstance(c, Entity) else c) for c in composed]
+
+
+def _on_deserialize_composed(refs):
+    return [r.chunk for r in refs]
+
+
+class FuseChunkData(ChunkData):
+
+    class ChunkRef(Serializable):
+        _chunk = OneOfField('chunk', tensor_chunk='mars.tensor.core.TensorChunkData',
+                            tensor='mars.tensor.core.TensorData',
+                            dataframe_chunk='mars.dataframe.core.DataFrameChunkData',
+                            dataframe='mars.dataframe.core.DataFrameData',
+                            index_chunk='mars.dataframe.core.IndexChunkData',
+                            index='mars.dataframe.core.IndexData',
+                            series_chunk='mars.dataframe.core.SeriesChunkData',
+                            series='mars.dataframe.core.SeriesData')
+
+        @property
+        def chunk(self):
+            return self._chunk
+
+    _composed = ListField('composed', ValueType.reference(ChunkRef),
+                          on_serialize=_on_serialize_composed,
+                          on_deserialize=_on_deserialize_composed)
+
+    @classmethod
+    def cls(cls, provider):
+        if provider.type == ProviderType.protobuf:
+            from .serialize.protos.fusechunk_pb2 import FuseChunkDef
+            return FuseChunkDef
+        return super(FuseChunkData, cls).cls(provider)
+
+    @property
+    def params(self):
+        # params return the properties which useful to rebuild a new chunk
+        p = {
+            'index': self.index,
+        }
+        if self.dtype:
+            p['dtype'] = self.dtype
+        return p
+
+    @property
+    def shape(self):
+        return self._extra_params.get('shape')
+
+    @property
+    def dtype(self):
+        # have dtype when the last compose is tensor
+        return self._extra_params.get('dtype') or self.op.dtype
+
+    @property
+    def nbytes(self):
+        return np.prod(self.shape) * self.dtype.itemsize
+
+
+class FuseChunk(Chunk):
+    __slots__ = ()
+    _allow_data_type_ = (FuseChunkData,)
+
+
+FUSE_CHUNK_TYPE = (FuseChunkData, FuseChunk)
+
+
 class TileableData(SerializableWithKey, Tileable):
     __slots__ = '__weakref__', '_siblings', '_cix'
     _no_copy_attrs_ = SerializableWithKey._no_copy_attrs_ | {'_cix'}

mars/dataframe/core.py

@@ -260,6 +260,13 @@ class IndexChunkData(ChunkData):
     _dtype = DataTypeField('dtype')
     _index_value = ReferenceField('index_value', IndexValue)
 
+    @classmethod
+    def cls(cls, provider):
+        if provider.type == ProviderType.protobuf:
+            from ..serialize.protos.dataframe_pb2 import IndexChunkDef
+            return IndexChunkDef
+        return super(IndexChunkData, cls).cls(provider)
+
     @property
     def dtype(self):
         return self._dtype
@@ -314,6 +321,13 @@ class SeriesChunkData(ChunkData):
     _dtype = DataTypeField('dtype')
     _index_value = ReferenceField('index_value', IndexValue)
 
+    @classmethod
+    def cls(cls, provider):
+        if provider.type == ProviderType.protobuf:
+            from ..serialize.protos.dataframe_pb2 import SeriesChunkDef
+            return SeriesChunkDef
+        return super(SeriesChunkData, cls).cls(provider)
+
     @property
     def dtype(self):
         return self._dtype
@@ -375,9 +389,16 @@ class DataFrameChunkData(ChunkData):
     _index_value = ReferenceField('index_value', IndexValue)
     _columns_value = ReferenceField('columns_value', IndexValue)
 
+    @classmethod
+    def cls(cls, provider):
+        if provider.type == ProviderType.protobuf:
+            from ..serialize.protos.dataframe_pb2 import DataFrameChunkDef
+            return DataFrameChunkDef
+        return super(DataFrameChunkData, cls).cls(provider)
+
     @property
     def params(self):
-        # params return the properties which useful to rebuild a new tileable object
+        # params return the properties which useful to rebuild a new chunk
         return {
             'shape': self.shape,
             'dtypes': self.dtypes,

mars/graph.pyx

@@ -320,6 +320,7 @@ cdef class DirectedGraph:
 
         from .tensor.core import CHUNK_TYPE as TENSOR_CHUNK_TYPE, TENSOR_TYPE, Chunk, Tensor
         from .dataframe.core import CHUNK_TYPE as DATAFRAME_CHUNK_TYPE, DATAFRAME_TYPE, DataFrame
+        from .core import FUSE_CHUNK_TYPE
 
         level = None
 
@@ -334,7 +335,7 @@ cdef class DirectedGraph:
                 visited.add(c)
 
         for node in self.iter_nodes():
-            if isinstance(node, TENSOR_CHUNK_TYPE + DATAFRAME_CHUNK_TYPE):
+            if isinstance(node, TENSOR_CHUNK_TYPE + DATAFRAME_CHUNK_TYPE + FUSE_CHUNK_TYPE):
                 node = node.data if isinstance(node, Chunk) else node
                 add_obj(node)
                 if node.composed:
@@ -510,7 +511,12 @@ class SerializableGraphNode(Serializable):
                        tensor_chunk='mars.tensor.core.TensorChunkData',
                        tensor='mars.tensor.core.TensorData',
                        dataframe_chunk='mars.dataframe.core.DataFrameChunkData',
-                       dataframe='mars.dataframe.core.DataFrameData')
+                       dataframe='mars.dataframe.core.DataFrameData',
+                       index_chunk='mars.dataframe.core.IndexChunkData',
+                       index='mars.dataframe.core.Index',
+                       series_chunk='mars.dataframe.core.SeriesChunkData',
+                       series='mars.dataframe.core.Series',
+                       fuse_chunk='mars.core.FuseChunkData')
 
     @classmethod
     def cls(cls, provider):

mars/serialize/protos/chunk.proto

@@ -2,7 +2,18 @@ syntax = "proto3";
 
 import "mars/serialize/protos/value.proto";
 import "mars/serialize/protos/indexvalue.proto";
-import "mars/serialize/protos/chunk.proto";
+
+
+message IndexChunkDef {
+    string key = 1;
+    repeated uint32 index = 2 [packed = true];
+    repeated int64 shape = 3;
+    Value op = 4;  // just store key here
+    bool cached = 5;
+    Value dtype = 6;
+    Value extra_params = 8;
+    string id = 9;
+}
 
 
 message IndexDef {
@@ -11,36 +22,64 @@ message IndexDef {
     Value dtype = 3;
     Value op = 4; // store operand's key and id
     Value nsplits = 5;
-    repeated ChunkDef chunks = 6;
+    repeated IndexChunkDef chunks = 6;
     Value extra_params = 7;
     string id = 8;
     IndexValue index_value = 9;
 }
 
 
+message SeriesChunkDef {
+    string key = 1;
+    repeated uint32 index = 2 [packed = true];
+    repeated int64 shape = 3;
+    Value op = 4;  // just store key here
+    bool cached = 5;
+    Value dtype = 6;
+    IndexValue index_value = 11;
+    Value extra_params = 8;
+    string id = 9;
+}
+
+
 message SeriesDef {
     string key = 1;
     repeated int64 shape = 2;
     Value dtype = 3;
     Value op = 4;  // store operand's key and id
     Value nsplits = 5;
-    repeated ChunkDef chunks = 6;
+    repeated IndexChunkDef chunks = 6;
     Value extra_params = 7;
     string id = 8;
     Value name = 9;
     IndexValue index_value = 10;
 }
 
 
+message DataFrameChunkDef {
+    string key = 1;
+    repeated uint32 index = 2 [packed = true];
+    repeated int64 shape = 3;
+    Value op = 4;  // just store key here
+    bool cached = 5;
+    Value dtypes = 10;
+    IndexValue index_value = 11;
+    IndexValue columns_value = 12;
+    Value extra_params = 8;
+    string id = 9;
+}
+
+
 message DataFrameDef {
     string key = 1;
     repeated int64 shape = 2;
     Value dtypes = 3;
     Value op = 4;  // store operand's key and id
     Value nsplits = 5;
-    repeated ChunkDef chunks = 6;
+    repeated DataFrameChunkDef chunks = 6;
     Value extra_params = 7;
     string id = 8;
     IndexValue index_value = 9;
     IndexValue columns_value = 10;
 }
+