Specialize Serialization for OrderedDict (#1035)

Specialize Serialization for OrderedDict and defaultdict
This commit is contained in:
Philipp Moritz
2017-10-02 17:33:10 -07:00
committed by Robert Nishihara
parent b94d85fb5d
commit 57bd1d6ff5
3 changed files with 36 additions and 6 deletions
+22
View File
@@ -1053,6 +1053,28 @@ def _initialize_serialization(worker=global_worker):
custom_serializer=array_custom_serializer,
custom_deserializer=array_custom_deserializer)
def ordered_dict_custom_serializer(obj):
return list(obj.keys()), list(obj.values())
def ordered_dict_custom_deserializer(obj):
return collections.OrderedDict(zip(obj[0], obj[1]))
worker.serialization_context.register_type(
collections.OrderedDict, 20 * b"\x02", pickle=False,
custom_serializer=ordered_dict_custom_serializer,
custom_deserializer=ordered_dict_custom_deserializer)
def default_dict_custom_serializer(obj):
return list(obj.keys()), list(obj.values()), obj.default_factory
def default_dict_custom_deserializer(obj):
return collections.defaultdict(obj[2], zip(obj[0], obj[1]))
worker.serialization_context.register_type(
collections.defaultdict, 20 * b"\x03", pickle=False,
custom_serializer=default_dict_custom_serializer,
custom_deserializer=default_dict_custom_deserializer)
if worker.mode in [SCRIPT_MODE, SILENT_MODE]:
# These should only be called on the driver because _register_class
# will export the class to all of the workers.