[Data] Support PyArrow 14.0.1

[bveeramani]

Why are these changes needed?

Arrow nightly tests are failing because our custom tensor extension uses pyarrow.PyExtensionType, and apache/arrow#38608 deprecated pyarrow.PyExtensionType in favor of pyarrow.ExtensionType. This PR updates our tensor extension to use the now-recommended APIs.

https://buildkite.com/ray-project/postmerge/builds/1593#018baae8-726d-451c-b4e4-db9e95113b81

Related issue number

Fixes #41139

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[bveeramani]

Will uncomment this after CI passes

[bveeramani]

Closing for now. Team can pick up when we decide to prioritize #41139

[pitrou]

In case you are not aware, this is making PyArrow vulnerable to https://www.cve.org/CVERecord?id=CVE-2023-47248

[jorisvandenbossche]

I understand that you want to keep it working to read files created by older versions of Ray (because you know those extension types are safe), but the problem with this general enabling of PyExtensionType is that just having ray imported opens the security hole for your users also when reading other files.

A possible alternative would be to register a PyExtensionType class that uses a custom unpickler that only allows to unpickle your own type, and raises an error for anything else.
A custom unpickler could look like:

import io
import pickle

class RestrictedUnpickler(pickle.Unpickler):

    def find_class(self, module, name):
        if (module, name) == ("ray.air.util.tensor_extensions.arrow", "ArrowTensorType"):
            from ray.air.util.tensor_extensions.arrow import ArrowTensorType
            return ArrowTensorType
        elif (module, name) == ("pyarrow.lib", "type_for_alias"):
            from pyarrow import type_for_alias
            return type_for_alias
        else:
            raise pickle.UnpicklingError("loading '%s.%s' is forbidden" % (module, name))

RestrictedUnpickler(io.BytesIO(serialized)).load()

That was my quick attempt for ArrowTensorType based on limited knowledge of the code base; it might be too restrictive in general (I don't know which pyarrow data types are allowed). But it should be possible to create something like this that works for all your extension types while still restricting any random pickled data.

[bveeramani]

@jorisvandenbossche thanks for the guidance. How would we make PyArrow use the custom unpickler?

[pitrou]

I think you can monkeypatch PyArrow just like the hotfix does:
https://github.com/pitrou/pyarrow-hotfix/blob/69b94597e6f96707b7a7782468fac0920b1ad74a/src/pyarrow_hotfix/__init__.py#L64-L68

[bveeramani]

@pitrou and then we'd use the custom unpickler in __arrow_ext_deserialize__?

[pitrou]

Yes, exactly.

[bveeramani]

Gotcha! Thanks!

[jorisvandenbossche]

You can also take a look at pyarrow's own implementation of PyExtensionType (https://github.com/apache/arrow/blob/43f4c286fef92b709853eefef207a50aefd9f5d3/python/pyarrow/types.pxi#L1700-L1742), and then re-register your own version like the pyarrow-hotfix does.