Repository: dask/dask
File: dask/tokenize.py (Python)

Code before:
    import numba

    @normalize_token.register(numba.core.serialize.ReduceMixin)
    def normalize_numba_ufunc(obj):
        return normalize_token((obj._reduce_class(), obj._reduce_states()))


@normalize_token.register_lazy("pyarrow")
def register_pyarrow():
    import pyarrow as pa

    @normalize_token.register(pa.DataType)
    def normalize_datatype(dt):
        return pickle.dumps(dt, protocol=4)

    @normalize_token.register(pa.Table)
    def normalize_table(dt):
        return (
            "pa.Table",
            normalize_token(dt.schema),
            normalize_token(dt.columns),
        )

    @normalize_token.register(pa.ChunkedArray)
    def normalize_chunked_array(arr):
        # So, pyarrow chooses to consolidate chunks
        # when pickling. To get deterministic tokenization, we need to
        # do the same thing...
        # I really don't want to do that this. This is a potentially
        # expensive operation.
        return ("pa.ChunkedArray", normalize_token(arr.combine_chunks()))
        # return (
        #     "pa.ChunkedArray",
        #     normalize_token(arr.type),
        #     normalize_token(arr.chunks),
        # )

    @normalize_token.register(pa.Array)
    def normalize_chunked_array(arr):
        buffers = arr.buffers()
        # pyarrow does something clever when (de)serializing an array that has
        # an empty validity map: The buffers for the deserialized array will
        # have `None` instead of the empty validity map.
        #
        # We'll replicate that behavior here to ensure we get consistent
        # tokenization.
        buffers = arr.buffers()
        if len(buffers) and buffers[0] is not None and arr.null_count == 0:
            buffers[0] = None

        return (

Review comment (performance):
In the tests, we assert that `tokenize(thing)` equals `tokenize(cloudpickle.loads(cloudpickle.dumps(thing)))`.

I need to dig into what pyarrow is doing when pickling Arrays / ChunkedArrays, but it seems to "normalize" the data in a couple of ways

1. ChunkedArrays with multiple chunks are combined into a single chunk (this change)
2. ChunkedArrays with a validity map but zero nulls do not allocate a validity map when round-tripped through pickle.

I don't care so much about the validity map (it's cheap to check and discard). But calling `combine_chunks()` here is I think relatively expensive since we may need to memcpy the buffers into a new location.

Code after:

    @normalize_token.register(pd.api.extensions.ExtensionDtype)
    def normalize_period_dtype(dtype):
        return normalize_token(dtype.name)

    @normalize_token.register(type(pd.NA))
    def normalize_na(na):
        return pd.NA

    @normalize_token.register(pd.offsets.BaseOffset)
    def normalize_offset(offset):
        return offset.freqstr


@normalize_token.register_lazy("numba")
def register_numba():
    import numba

    @normalize_token.register(numba.core.serialize.ReduceMixin)
    def normalize_numba_ufunc(obj):
        return normalize_token((obj._reduce_class(), obj._reduce_states()))


@normalize_token.register_lazy("pyarrow")
def register_pyarrow():
    import pyarrow as pa

    @normalize_token.register(pa.DataType)
    def normalize_datatype(dt):
        return pickle.dumps(dt, protocol=4)

    @normalize_token.register(pa.Table)
    def normalize_table(dt):
        return (
            "pa.Table",
            normalize_token(dt.schema),
            normalize_token(dt.columns),
        )

    @normalize_token.register(pa.ChunkedArray)
    def normalize_chunked_array(arr):
        return (
            "pa.ChunkedArray",
            normalize_token(arr.type),
            normalize_token(arr.chunks),
        )

    @normalize_token.register(pa.Array)
    def normalize_array(arr):
        buffers = arr.buffers()
        # pyarrow does something clever when (de)serializing an array that has