[
https://issues.apache.org/jira/browse/ARROW-2024?page=com.atlassian.jira.plugin.system.issuetabpanels:comment-tabpanel&focusedCommentId=16348940#comment-16348940
]
ASF GitHub Bot commented on ARROW-2024:
---------------------------------------
wesm closed pull request #1538: ARROW-2024: [Python] Remove torch serialization
from default serialization context.
URL: https://github.com/apache/arrow/pull/1538
This is a PR merged from a forked repository.
As GitHub hides the original diff on merge, it is displayed below for
the sake of provenance:
As this is a foreign pull request (from a fork), the diff is supplied
below (as it won't show otherwise due to GitHub magic):
diff --git a/python/pyarrow/__init__.py b/python/pyarrow/__init__.py
index a245fe679..8b3cba924 100644
--- a/python/pyarrow/__init__.py
+++ b/python/pyarrow/__init__.py
@@ -124,9 +124,10 @@
localfs = LocalFileSystem.get_instance()
-from pyarrow.serialization import (_default_serialization_context,
+from pyarrow.serialization import (default_serialization_context,
pandas_serialization_context,
- register_default_serialization_handlers)
+ register_default_serialization_handlers,
+ register_torch_serialization_handlers)
import pyarrow.types as types
diff --git a/python/pyarrow/serialization.py b/python/pyarrow/serialization.py
index 61f2e83f3..c8b72b748 100644
--- a/python/pyarrow/serialization.py
+++ b/python/pyarrow/serialization.py
@@ -22,7 +22,8 @@
import numpy as np
from pyarrow.compat import builtin_pickle
-from pyarrow.lib import _default_serialization_context, frombuffer
+from pyarrow.lib import (SerializationContext, _default_serialization_context,
+ frombuffer)
try:
import cloudpickle
@@ -102,6 +103,31 @@ def _deserialize_pandas_series(data):
custom_deserializer=_deserialize_pandas_dataframe)
+def register_torch_serialization_handlers(serialization_context):
+ # ----------------------------------------------------------------------
+ # Set up serialization for pytorch tensors
+
+ try:
+ import torch
+
+ def _serialize_torch_tensor(obj):
+ return obj.numpy()
+
+ def _deserialize_torch_tensor(data):
+ return torch.from_numpy(data)
+
+ for t in [torch.FloatTensor, torch.DoubleTensor, torch.HalfTensor,
+ torch.ByteTensor, torch.CharTensor, torch.ShortTensor,
+ torch.IntTensor, torch.LongTensor]:
+ serialization_context.register_type(
+ t, "torch." + t.__name__,
+ custom_serializer=_serialize_torch_tensor,
+ custom_deserializer=_deserialize_torch_tensor)
+ except ImportError:
+ # no torch
+ pass
+
+
def register_default_serialization_handlers(serialization_context):
# ----------------------------------------------------------------------
@@ -154,37 +180,21 @@ def _deserialize_default_dict(data):
custom_serializer=_serialize_numpy_array_list,
custom_deserializer=_deserialize_numpy_array_list)
- # ----------------------------------------------------------------------
- # Set up serialization for pytorch tensors
-
- try:
- import torch
-
- def _serialize_torch_tensor(obj):
- return obj.numpy()
+ _register_custom_pandas_handlers(serialization_context)
- def _deserialize_torch_tensor(data):
- return torch.from_numpy(data)
- for t in [torch.FloatTensor, torch.DoubleTensor, torch.HalfTensor,
- torch.ByteTensor, torch.CharTensor, torch.ShortTensor,
- torch.IntTensor, torch.LongTensor]:
- serialization_context.register_type(
- t, "torch." + t.__name__,
- custom_serializer=_serialize_torch_tensor,
- custom_deserializer=_deserialize_torch_tensor)
- except ImportError:
- # no torch
- pass
-
- _register_custom_pandas_handlers(serialization_context)
+def default_serialization_context():
+ context = SerializationContext()
+ register_default_serialization_handlers(context)
+ return context
register_default_serialization_handlers(_default_serialization_context)
-pandas_serialization_context = _default_serialization_context.clone()
-pandas_serialization_context.register_type(
- np.ndarray, 'np.array',
- custom_serializer=_serialize_numpy_array_pickle,
- custom_deserializer=_deserialize_numpy_array_pickle)
+def pandas_serialization_context():
+ context = default_serialization_context()
+ context.register_type(np.ndarray, 'np.array',
+ custom_serializer=_serialize_numpy_array_pickle,
+ custom_deserializer=_deserialize_numpy_array_pickle)
+ return context
diff --git a/python/pyarrow/tests/test_convert_pandas.py
b/python/pyarrow/tests/test_convert_pandas.py
index fa265e55c..ca2f1e361 100644
--- a/python/pyarrow/tests/test_convert_pandas.py
+++ b/python/pyarrow/tests/test_convert_pandas.py
@@ -1404,7 +1404,7 @@ def _fully_loaded_dataframe_example():
def _check_serialize_components_roundtrip(df):
- ctx = pa.pandas_serialization_context
+ ctx = pa.pandas_serialization_context()
components = ctx.serialize(df).to_components()
deserialized = ctx.deserialize_components(components)
diff --git a/python/pyarrow/tests/test_serialization.py
b/python/pyarrow/tests/test_serialization.py
index 284c7fb4c..7a420106f 100644
--- a/python/pyarrow/tests/test_serialization.py
+++ b/python/pyarrow/tests/test_serialization.py
@@ -190,8 +190,7 @@ class CustomError(Exception):
def make_serialization_context():
-
- context = pa._default_serialization_context
+ context = pa.default_serialization_context()
context.register_type(Foo, "Foo")
context.register_type(Bar, "Bar")
@@ -207,26 +206,27 @@ def make_serialization_context():
return context
-serialization_context = make_serialization_context()
+global_serialization_context = make_serialization_context()
-def serialization_roundtrip(value, scratch_buffer, ctx=serialization_context):
+def serialization_roundtrip(value, scratch_buffer,
+ context=global_serialization_context):
writer = pa.FixedSizeBufferWriter(scratch_buffer)
- pa.serialize_to(value, writer, ctx)
+ pa.serialize_to(value, writer, context=context)
reader = pa.BufferReader(scratch_buffer)
- result = pa.deserialize_from(reader, None, ctx)
+ result = pa.deserialize_from(reader, None, context=context)
assert_equal(value, result)
- _check_component_roundtrip(value)
+ _check_component_roundtrip(value, context=context)
-def _check_component_roundtrip(value):
+def _check_component_roundtrip(value, context=global_serialization_context):
# Test to/from components
- serialized = pa.serialize(value)
+ serialized = pa.serialize(value, context=context)
components = serialized.to_components()
from_comp = pa.SerializedPyObject.from_components(components)
- recons = from_comp.deserialize()
+ recons = from_comp.deserialize(context=context)
assert_equal(value, recons)
@@ -252,7 +252,7 @@ def test_primitive_serialization(large_buffer):
for obj in PRIMITIVE_OBJECTS:
serialization_roundtrip(obj, large_buffer)
serialization_roundtrip(obj, large_buffer,
- pa.pandas_serialization_context)
+ pa.pandas_serialization_context())
def test_serialize_to_buffer():
@@ -318,22 +318,26 @@ def test_datetime_serialization(large_buffer):
def test_torch_serialization(large_buffer):
pytest.importorskip("torch")
import torch
+
+ serialization_context = pa.default_serialization_context()
+ pa.register_torch_serialization_handlers(serialization_context)
# These are the only types that are supported for the
# PyTorch to NumPy conversion
for t in ["float32", "float64",
"uint8", "int16", "int32", "int64"]:
obj = torch.from_numpy(np.random.randn(1000).astype(t))
- serialization_roundtrip(obj, large_buffer)
+ serialization_roundtrip(obj, large_buffer,
+ context=serialization_context)
def test_numpy_immutable(large_buffer):
obj = np.zeros([10])
writer = pa.FixedSizeBufferWriter(large_buffer)
- pa.serialize_to(obj, writer, serialization_context)
+ pa.serialize_to(obj, writer, global_serialization_context)
reader = pa.BufferReader(large_buffer)
- result = pa.deserialize_from(reader, None, serialization_context)
+ result = pa.deserialize_from(reader, None, global_serialization_context)
with pytest.raises(ValueError):
result[0] = 1.0
@@ -351,12 +355,12 @@ def serialize_dummy_class(obj):
def deserialize_dummy_class(serialized_obj):
return serialized_obj
- pa._default_serialization_context.register_type(
- DummyClass, "DummyClass",
- custom_serializer=serialize_dummy_class,
- custom_deserializer=deserialize_dummy_class)
+ context = pa.default_serialization_context()
+ context.register_type(DummyClass, "DummyClass",
+ custom_serializer=serialize_dummy_class,
+ custom_deserializer=deserialize_dummy_class)
- pa.serialize(DummyClass())
+ pa.serialize(DummyClass(), context=context)
def test_buffer_serialization():
@@ -370,13 +374,14 @@ def serialize_buffer_class(obj):
def deserialize_buffer_class(serialized_obj):
return serialized_obj
- pa._default_serialization_context.register_type(
+ context = pa.default_serialization_context()
+ context.register_type(
BufferClass, "BufferClass",
custom_serializer=serialize_buffer_class,
custom_deserializer=deserialize_buffer_class)
- b = pa.serialize(BufferClass()).to_buffer()
- assert pa.deserialize(b).to_pybytes() == b"hello"
+ b = pa.serialize(BufferClass(), context=context).to_buffer()
+ assert pa.deserialize(b, context=context).to_pybytes() == b"hello"
@pytest.mark.skip(reason="extensive memory requirements")
@@ -485,15 +490,16 @@ def test_serialize_subclasses():
# with register_type will result in faster and more memory
# efficient serialization.
- serialization_context.register_type(
+ context = pa.default_serialization_context()
+ context.register_type(
Serializable, "Serializable",
custom_serializer=serialize_serializable,
custom_deserializer=deserialize_serializable)
a = SerializableClass()
- serialized = pa.serialize(a)
+ serialized = pa.serialize(a, context=context)
- deserialized = serialized.deserialize()
+ deserialized = serialized.deserialize(context=context)
assert type(deserialized).__name__ == SerializableClass.__name__
assert deserialized.value == 3
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
> Remove global SerializationContext variables.
> ---------------------------------------------
>
> Key: ARROW-2024
> URL: https://issues.apache.org/jira/browse/ARROW-2024
> Project: Apache Arrow
> Issue Type: Improvement
> Components: Python
> Reporter: Robert Nishihara
> Priority: Minor
> Labels: pull-request-available
> Fix For: 0.9.0
>
>
> We should get rid of the global variables
> {{{{_default_serialization_context}}}} and
> {{{{pandas_serialization_context}}}} and replace them with functions
> {{{{default_serialization_context()}}}} and
> {{{{pandas_serialization_context()}}}}.
> This will also make it faster to do {{{{import pyarrow}}}}.
--
This message was sent by Atlassian JIRA
(v7.6.3#76005)
