ray-project--ray/python/ray/_private/arrow_serialization.py

# arrow_serialization.py must resides outside of ray.data, otherwise
# it causes circular dependency issues for AsyncActors due to
# ray.data's lazy import.
# see https://github.com/ray-project/ray/issues/30498 for more context.
import logging
import os
import sys
from dataclasses import dataclass
from typing import TYPE_CHECKING, Callable, List, Optional, Tuple

from ray._private.utils import is_in_test

if TYPE_CHECKING:
    import pyarrow

RAY_DISABLE_CUSTOM_ARROW_JSON_OPTIONS_SERIALIZATION = (
    "RAY_DISABLE_CUSTOM_ARROW_JSON_OPTIONS_SERIALIZATION"
)
RAY_DISABLE_CUSTOM_ARROW_DATA_SERIALIZATION = (
    "RAY_DISABLE_CUSTOM_ARROW_DATA_SERIALIZATION"
)

logger = logging.getLogger(__name__)

# Whether we have already warned the user about bloated fallback serialization.
_serialization_fallback_set = set()


def _register_custom_datasets_serializers(serialization_context):
    try:
        import pyarrow as pa  # noqa: F401
    except ModuleNotFoundError:
        # No pyarrow installed so not using Arrow, so no need for custom serializers.
        return

    # Register all custom serializers required by Datasets.
    _register_arrow_data_serializer(serialization_context)
    _register_arrow_json_readoptions_serializer(serialization_context)
    _register_arrow_json_parseoptions_serializer(serialization_context)


# Register custom Arrow JSON ReadOptions serializer to workaround it not being picklable
# in Arrow < 8.0.0.
def _register_arrow_json_readoptions_serializer(serialization_context):
    if (
        os.environ.get(
            RAY_DISABLE_CUSTOM_ARROW_JSON_OPTIONS_SERIALIZATION,
            "0",
        )
        == "1"
    ):
        return

    import pyarrow.json as pajson

    serialization_context._register_cloudpickle_serializer(
        pajson.ReadOptions,
        custom_serializer=lambda opts: (opts.use_threads, opts.block_size),
        custom_deserializer=lambda args: pajson.ReadOptions(*args),
    )


def _register_arrow_json_parseoptions_serializer(serialization_context):
    if (
        os.environ.get(
            RAY_DISABLE_CUSTOM_ARROW_JSON_OPTIONS_SERIALIZATION,
            "0",
        )
        == "1"
    ):
        return

    import pyarrow.json as pajson

    serialization_context._register_cloudpickle_serializer(
        pajson.ParseOptions,
        custom_serializer=lambda opts: (
            opts.explicit_schema,
            opts.newlines_in_values,
            opts.unexpected_field_behavior,
        ),
        custom_deserializer=lambda args: pajson.ParseOptions(*args),
    )


# Register custom Arrow data serializer to work around zero-copy slice pickling bug.
# See https://issues.apache.org/jira/browse/ARROW-10739.
def _register_arrow_data_serializer(serialization_context):
    """Custom reducer for Arrow data that works around a zero-copy slicing pickling
    bug by using the Arrow IPC format for the underlying serialization.

    Background:
        Arrow has both array-level slicing and buffer-level slicing; both are zero-copy,
        but the former has a serialization bug where the entire buffer is serialized
        instead of just the slice, while the latter's serialization works as expected
        and only serializes the slice of the buffer. I.e., array-level slicing doesn't
        propagate the slice down to the buffer when serializing the array.

        We work around this by registering a custom cloudpickle reducers for Arrow
        Tables that delegates serialization to the Arrow IPC format; thankfully, Arrow's
        IPC serialization has fixed this buffer truncation bug.

    See https://issues.apache.org/jira/browse/ARROW-10739.
    """
    if os.environ.get(RAY_DISABLE_CUSTOM_ARROW_DATA_SERIALIZATION, "0") == "1":
        return

    import pyarrow as pa

    serialization_context._register_cloudpickle_reducer(pa.Table, _arrow_table_reduce)
    serialization_context._register_cloudpickle_reducer(pa.Schema, _arrow_schema_reduce)


def _arrow_schema_reduce(
    schema: "pyarrow.Schema",
) -> Tuple[Callable[["bytes"], "pyarrow.Schema"], Tuple[bytes]]:
    """Custom reducer for Arrow Schema that uses IPC serialization for performance.

    Arrow's native IPC serialization for schemas is significantly faster than
    cloudpickle (10-20x for serialization, 2-3x for deserialization), making
    this optimization particularly valuable for workloads with large schemas.
    """
    # Use Arrow's native IPC serialization which is much faster than cloudpickle
    return _restore_schema_from_ipc, (schema.serialize().to_pybytes(),)


def _restore_schema_from_ipc(buf: bytes) -> "pyarrow.Schema":
    """Restore an Arrow Schema serialized to Arrow IPC format."""
    import pyarrow as pa

    return pa.ipc.read_schema(pa.BufferReader(buf))


def _arrow_table_reduce(t: "pyarrow.Table"):
    """Custom reducer for Arrow Tables that works around a zero-copy slice pickling bug.
    Background:
        Arrow has both array-level slicing and buffer-level slicing; both are zero-copy,
        but the former has a serialization bug where the entire buffer is serialized
        instead of just the slice, while the latter's serialization works as expected
        and only serializes the slice of the buffer. I.e., array-level slicing doesn't
        propagate the slice down to the buffer when serializing the array.
        All that these copy methods do is, at serialization time, take the array-level
        slicing and translate them to buffer-level slicing, so only the buffer slice is
        sent over the wire instead of the entire buffer.
    See https://issues.apache.org/jira/browse/ARROW-10739.
    """
    global _serialization_fallback_set

    # Reduce the ChunkedArray columns.
    reduced_columns = []
    for column_name in t.column_names:
        column = t[column_name]
        try:
            # Delegate to ChunkedArray reducer.
            reduced_column = _arrow_chunked_array_reduce(column)
        except Exception as e:
            if not _is_dense_union(column.type) and is_in_test():
                # If running in a test and the column is not a dense union array
                # (which we expect to need a fallback), we want to raise the error,
                # not fall back.
                raise e from None
            if type(column.type) not in _serialization_fallback_set:
                logger.warning(
                    "Failed to complete optimized serialization of Arrow Table, "
                    f"serialization of column '{column_name}' of type {column.type} "
                    "failed, so we're falling back to Arrow IPC serialization for the "
                    "table. Note that this may result in slower serialization and more "
                    "worker memory utilization. Serialization error:",
                    exc_info=True,
                )
                _serialization_fallback_set.add(type(column.type))
            # Fall back to Arrow IPC-based workaround for the entire table.
            return _arrow_table_ipc_reduce(t)
        else:
            # Column reducer succeeded, add reduced column to list.
            reduced_columns.append(reduced_column)
    return _reconstruct_table, (reduced_columns, t.schema)


def _reconstruct_table(
    reduced_columns: List[Tuple[List["pyarrow.Array"], "pyarrow.DataType"]],
    schema: "pyarrow.Schema",
) -> "pyarrow.Table":
    """Restore a serialized Arrow Table, reconstructing each reduced column."""
    import pyarrow as pa

    # Reconstruct each reduced column.
    columns = []
    for chunks_payload, type_ in reduced_columns:
        columns.append(_reconstruct_chunked_array(chunks_payload, type_))

    return pa.Table.from_arrays(columns, schema=schema)


def _arrow_chunked_array_reduce(
    ca: "pyarrow.ChunkedArray",
) -> Tuple[List["PicklableArrayPayload"], "pyarrow.DataType"]:
    """Custom reducer for Arrow ChunkedArrays that works around a zero-copy slice
    pickling bug. This reducer does not return a reconstruction function, since it's
    expected to be reconstructed by the Arrow Table reconstructor.
    """
    # Convert chunks to serialization payloads.
    chunk_payloads = []
    for chunk in ca.chunks:
        chunk_payload = PicklableArrayPayload.from_array(chunk)
        chunk_payloads.append(chunk_payload)
    return chunk_payloads, ca.type


def _reconstruct_chunked_array(
    chunks: List["PicklableArrayPayload"], type_: "pyarrow.DataType"
) -> "pyarrow.ChunkedArray":
    """Restore a serialized Arrow ChunkedArray from chunks and type."""
    import pyarrow as pa

    # Reconstruct chunks from serialization payloads.
    chunks = [chunk.to_array() for chunk in chunks]
    return pa.chunked_array(chunks, type_)


@dataclass
class PicklableArrayPayload:
    """Picklable array payload, holding data buffers and array metadata.

    This is a helper container for pickling and reconstructing nested Arrow Arrays while
    ensuring that the buffers that underly zero-copy slice views are properly truncated.
    """

    # Array type.
    type: "pyarrow.DataType"
    # Length of array.
    length: int
    # Underlying data buffers.
    buffers: List["pyarrow.Buffer"]
    # Cached null count.
    null_count: int
    # Slice offset into base array.
    offset: int
    # Serialized array payloads for nested (child) arrays.
    children: List["PicklableArrayPayload"]

    @classmethod
    def from_array(self, a: "pyarrow.Array") -> "PicklableArrayPayload":
        """Create a picklable array payload from an Arrow Array.

        This will recursively accumulate data buffer and metadata payloads that are
        ready for pickling; namely, the data buffers underlying zero-copy slice views
        will be properly truncated.
        """
        return _array_to_array_payload(a)

    def to_array(self) -> "pyarrow.Array":
        """Reconstruct an Arrow Array from this picklable payload."""
        return _array_payload_to_array(self)


def _array_payload_to_array(payload: "PicklableArrayPayload") -> "pyarrow.Array":
    """Reconstruct an Arrow Array from a possibly nested PicklableArrayPayload."""
    import pyarrow as pa

    children = [child_payload.to_array() for child_payload in payload.children]

    if pa.types.is_dictionary(payload.type):
        # Dedicated path for reconstructing a DictionaryArray, since
        # Array.from_buffers() doesn't work for DictionaryArrays.
        assert len(children) == 2, len(children)
        indices, dictionary = children
        return pa.DictionaryArray.from_arrays(
            indices,
            dictionary,
            ordered=payload.type.ordered,  # Explicitly pass the ordered flag to from_arrays() to prevent dropping it as ordered=False by default
        )
    elif pa.types.is_map(payload.type) and len(children) > 1:
        # In pyarrow<7.0.0, the underlying map child array is not exposed, so we work
        # with the key and item arrays.
        assert len(children) == 3, len(children)
        offsets, keys, items = children
        return pa.MapArray.from_arrays(offsets, keys, items)
    elif isinstance(payload.type, pa.BaseExtensionType):
        assert len(children) == 1, len(children)
        storage = children[0]
        return payload.type.wrap_array(storage)
    else:
        # Common case: use Array.from_buffers() to construct an array of a certain type.
        return pa.Array.from_buffers(
            type=payload.type,
            length=payload.length,
            buffers=payload.buffers,
            null_count=payload.null_count,
            offset=payload.offset,
            children=children,
        )


def _array_to_array_payload(a: "pyarrow.Array") -> "PicklableArrayPayload":
    """Serialize an Arrow Array to an PicklableArrayPayload for later pickling.

    This function's primary purpose is to dispatch to the handler for the input array
    type.
    """
    import pyarrow as pa

    if _is_dense_union(a.type):
        # Dense unions are not supported.
        # TODO(Clark): Support dense unions.
        raise NotImplementedError(
            "Custom slice view serialization of dense union arrays is not yet "
            "supported."
        )

    # Dispatch to handler for array type.
    if pa.types.is_null(a.type):
        return _null_array_to_array_payload(a)
    elif _is_primitive(a.type):
        return _primitive_array_to_array_payload(a)
    elif _is_binary(a.type):
        return _binary_array_to_array_payload(a)
    elif pa.types.is_list(a.type) or pa.types.is_large_list(a.type):
        return _list_array_to_array_payload(a)
    elif pa.types.is_fixed_size_list(a.type):
        return _fixed_size_list_array_to_array_payload(a)
    elif pa.types.is_struct(a.type):
        return _struct_array_to_array_payload(a)
    elif pa.types.is_union(a.type):
        return _union_array_to_array_payload(a)
    elif pa.types.is_dictionary(a.type):
        return _dictionary_array_to_array_payload(a)
    elif pa.types.is_map(a.type):
        return _map_array_to_array_payload(a)
    elif isinstance(a.type, pa.BaseExtensionType):
        return _extension_array_to_array_payload(a)
    else:
        raise ValueError("Unhandled Arrow array type:", a.type)


def _is_primitive(type_: "pyarrow.DataType") -> bool:
    """Whether the provided Array type is primitive (boolean, numeric, temporal or
    fixed-size binary)."""
    import pyarrow as pa

    return (
        pa.types.is_integer(type_)
        or pa.types.is_floating(type_)
        or pa.types.is_decimal(type_)
        or pa.types.is_boolean(type_)
        or pa.types.is_temporal(type_)
        or pa.types.is_fixed_size_binary(type_)
    )


def _is_binary(type_: "pyarrow.DataType") -> bool:
    """Whether the provided Array type is a variable-sized binary type."""
    import pyarrow as pa

    return (
        pa.types.is_string(type_)
        or pa.types.is_large_string(type_)
        or pa.types.is_binary(type_)
        or pa.types.is_large_binary(type_)
    )


def _null_array_to_array_payload(a: "pyarrow.NullArray") -> "PicklableArrayPayload":
    """Serialize null array to PicklableArrayPayload."""
    # Buffer scheme: [None]
    return PicklableArrayPayload(
        type=a.type,
        length=len(a),
        buffers=[None],  # Single null buffer is expected.
        null_count=a.null_count,
        offset=0,
        children=[],
    )


def _primitive_array_to_array_payload(a: "pyarrow.Array") -> "PicklableArrayPayload":
    """Serialize primitive (numeric, temporal, boolean) arrays to
    PicklableArrayPayload.
    """
    assert _is_primitive(a.type), a.type
    # Buffer scheme: [bitmap, data]
    buffers = a.buffers()
    assert len(buffers) == 2, len(buffers)

    # Copy bitmap buffer, if needed.
    bitmap_buf = buffers[0]
    if a.null_count > 0:
        bitmap_buf = _copy_bitpacked_buffer_if_needed(bitmap_buf, a.offset, len(a))
    else:
        bitmap_buf = None

    # Copy data buffer, if needed.
    data_buf = buffers[1]
    if data_buf is not None:
        data_buf = _copy_buffer_if_needed(buffers[1], a.type, a.offset, len(a))

    return PicklableArrayPayload(
        type=a.type,
        length=len(a),
        buffers=[bitmap_buf, data_buf],
        null_count=a.null_count,
        offset=0,
        children=[],
    )


def _binary_array_to_array_payload(a: "pyarrow.Array") -> "PicklableArrayPayload":
    """Serialize binary (variable-sized binary, string) arrays to
    PicklableArrayPayload.
    """
    assert _is_binary(a.type), a.type
    # Buffer scheme: [bitmap, value_offsets, data]
    buffers = a.buffers()
    assert len(buffers) == 3, len(buffers)

    # Copy bitmap buffer, if needed.
    if a.null_count > 0:
        bitmap_buf = _copy_bitpacked_buffer_if_needed(buffers[0], a.offset, len(a))
    else:
        bitmap_buf = None

    # Copy offset buffer, if needed.
    offset_buf = buffers[1]
    offset_buf, data_offset, data_length = _copy_offsets_buffer_if_needed(
        offset_buf, a.type, a.offset, len(a)
    )
    data_buf = buffers[2]
    data_buf = _copy_buffer_if_needed(data_buf, None, data_offset, data_length)
    return PicklableArrayPayload(
        type=a.type,
        length=len(a),
        buffers=[bitmap_buf, offset_buf, data_buf],
        null_count=a.null_count,
        offset=0,
        children=[],
    )


def _list_array_to_array_payload(a: "pyarrow.Array") -> "PicklableArrayPayload":
    """Serialize list (regular and large) arrays to PicklableArrayPayload."""
    # Dedicated path for ListArrays. These arrays have a nested set of bitmap and
    # offset buffers, eventually bottoming out on a data buffer.
    # Buffer scheme:
    # [bitmap, offsets, bitmap, offsets, ..., bitmap, data]
    buffers = a.buffers()
    assert len(buffers) > 1, len(buffers)

    # Copy bitmap buffer, if needed.
    if a.null_count > 0:
        bitmap_buf = _copy_bitpacked_buffer_if_needed(buffers[0], a.offset, len(a))
    else:
        bitmap_buf = None

    # Copy offset buffer, if needed.
    offset_buf = buffers[1]
    offset_buf, child_offset, child_length = _copy_offsets_buffer_if_needed(
        offset_buf, a.type, a.offset, len(a)
    )

    # Propagate slice to child.
    child = a.values.slice(child_offset, child_length)

    return PicklableArrayPayload(
        type=a.type,
        length=len(a),
        buffers=[bitmap_buf, offset_buf],
        null_count=a.null_count,
        offset=0,
        children=[_array_to_array_payload(child)],
    )


def _fixed_size_list_array_to_array_payload(
    a: "pyarrow.FixedSizeListArray",
) -> "PicklableArrayPayload":
    """Serialize fixed size list arrays to PicklableArrayPayload."""
    # Dedicated path for fixed-size lists.
    # Buffer scheme:
    # [bitmap, values_bitmap, values_data, values_subbuffers...]
    buffers = a.buffers()
    assert len(buffers) >= 1, len(buffers)

    # Copy bitmap buffer, if needed.
    if a.null_count > 0:
        bitmap_buf = _copy_bitpacked_buffer_if_needed(buffers[0], a.offset, len(a))
    else:
        bitmap_buf = None

    # Propagate slice to child.
    child_offset = a.type.list_size * a.offset
    child_length = a.type.list_size * len(a)
    child = a.values.slice(child_offset, child_length)

    return PicklableArrayPayload(
        type=a.type,
        length=len(a),
        buffers=[bitmap_buf],
        null_count=a.null_count,
        offset=0,
        children=[_array_to_array_payload(child)],
    )


def _struct_array_to_array_payload(a: "pyarrow.StructArray") -> "PicklableArrayPayload":
    """Serialize struct arrays to PicklableArrayPayload."""
    # Dedicated path for StructArrays.
    # StructArrays have a top-level bitmap buffer and one or more children arrays.
    # Buffer scheme: [bitmap, None, child_bitmap, child_data, ...]
    buffers = a.buffers()
    assert len(buffers) >= 1, len(buffers)

    # Copy bitmap buffer, if needed.
    if a.null_count > 0:
        bitmap_buf = _copy_bitpacked_buffer_if_needed(buffers[0], a.offset, len(a))
    else:
        bitmap_buf = None

    # Get field children payload.
    # Offsets and truncations are already propagated to the field arrays, so we can
    # serialize them as-is.
    children = [_array_to_array_payload(a.field(i)) for i in range(a.type.num_fields)]
    return PicklableArrayPayload(
        type=a.type,
        length=len(a),
        buffers=[bitmap_buf],
        null_count=a.null_count,
        offset=0,
        children=children,
    )


def _union_array_to_array_payload(a: "pyarrow.UnionArray") -> "PicklableArrayPayload":
    """Serialize union arrays to PicklableArrayPayload."""
    import pyarrow as pa

    # Dedicated path for UnionArrays.
    # UnionArrays have a top-level bitmap buffer and type code buffer, and one or
    # more children arrays.
    # Buffer scheme: [None, typecodes, child_bitmap, child_data, ...]
    assert not _is_dense_union(a.type)
    buffers = a.buffers()
    assert len(buffers) > 1, len(buffers)

    bitmap_buf = buffers[0]
    assert bitmap_buf is None, bitmap_buf

    # Copy type code buffer, if needed.
    type_code_buf = buffers[1]
    type_code_buf = _copy_buffer_if_needed(type_code_buf, pa.int8(), a.offset, len(a))

    # Get field children payload.
    # Offsets and truncations are already propagated to the field arrays, so we can
    # serialize them as-is.
    children = [_array_to_array_payload(a.field(i)) for i in range(a.type.num_fields)]
    return PicklableArrayPayload(
        type=a.type,
        length=len(a),
        buffers=[bitmap_buf, type_code_buf],
        null_count=a.null_count,
        offset=0,
        children=children,
    )


def _dictionary_array_to_array_payload(
    a: "pyarrow.DictionaryArray",
) -> "PicklableArrayPayload":
    """Serialize dictionary arrays to PicklableArrayPayload."""
    # Dedicated path for DictionaryArrays.
    # Buffer scheme: [indices_bitmap, indices_data] (dictionary stored separately)
    indices_payload = _array_to_array_payload(a.indices)
    dictionary_payload = _array_to_array_payload(a.dictionary)
    return PicklableArrayPayload(
        type=a.type,
        length=len(a),
        buffers=[],
        null_count=a.null_count,
        offset=0,
        children=[indices_payload, dictionary_payload],
    )


def _map_array_to_array_payload(a: "pyarrow.MapArray") -> "PicklableArrayPayload":
    """Serialize map arrays to PicklableArrayPayload."""
    import pyarrow as pa

    # Dedicated path for MapArrays.
    # Buffer scheme: [bitmap, offsets, child_struct_array_buffers, ...]
    buffers = a.buffers()
    assert len(buffers) > 0, len(buffers)

    # Copy bitmap buffer, if needed.
    if a.null_count > 0:
        bitmap_buf = _copy_bitpacked_buffer_if_needed(buffers[0], a.offset, len(a))
    else:
        bitmap_buf = None

    new_buffers = [bitmap_buf]

    # Copy offsets buffer, if needed.
    offset_buf = buffers[1]
    offset_buf, data_offset, data_length = _copy_offsets_buffer_if_needed(
        offset_buf, a.type, a.offset, len(a)
    )

    if isinstance(a, pa.lib.ListArray):
        # Map arrays directly expose the one child struct array in pyarrow>=7.0.0, which
        # is easier to work with than the raw buffers.
        new_buffers.append(offset_buf)
        children = [_array_to_array_payload(a.values.slice(data_offset, data_length))]
    else:
        # In pyarrow<7.0.0, the child struct array is not exposed, so we work with the
        # key and item arrays.
        buffers = a.buffers()
        assert len(buffers) > 2, len(buffers)
        # Reconstruct offsets array.
        offsets = pa.Array.from_buffers(
            pa.int32(), len(a) + 1, [bitmap_buf, offset_buf]
        )
        # Propagate slice to keys.
        keys = a.keys.slice(data_offset, data_length)
        # Propagate slice to items.
        items = a.items.slice(data_offset, data_length)
        children = [
            _array_to_array_payload(offsets),
            _array_to_array_payload(keys),
            _array_to_array_payload(items),
        ]
    return PicklableArrayPayload(
        type=a.type,
        length=len(a),
        buffers=new_buffers,
        null_count=a.null_count,
        offset=0,
        children=children,
    )


def _extension_array_to_array_payload(
    a: "pyarrow.ExtensionArray",
) -> "PicklableArrayPayload":
    storage_payload = _array_to_array_payload(a.storage)
    return PicklableArrayPayload(
        type=a.type,
        length=len(a),
        buffers=[],
        null_count=a.null_count,
        offset=0,
        children=[storage_payload],
    )


def _copy_buffer_if_needed(
    buf: "pyarrow.Buffer",
    type_: Optional["pyarrow.DataType"],
    offset: int,
    length: int,
) -> "pyarrow.Buffer":
    """Copy buffer, if needed."""
    import pyarrow as pa

    if type_ is not None and pa.types.is_boolean(type_):
        # Arrow boolean array buffers are bit-packed, with 8 entries per byte,
        # and are accessed via bit offsets.
        buf = _copy_bitpacked_buffer_if_needed(buf, offset, length)
    else:
        type_bytewidth = type_.bit_width // 8 if type_ is not None else 1
        buf = _copy_normal_buffer_if_needed(buf, type_bytewidth, offset, length)
    return buf


def _copy_normal_buffer_if_needed(
    buf: "pyarrow.Buffer",
    byte_width: int,
    offset: int,
    length: int,
) -> "pyarrow.Buffer":
    """Copy buffer, if needed."""
    byte_offset = offset * byte_width
    byte_length = length * byte_width
    if offset > 0 or byte_length < buf.size:
        # Array is a zero-copy slice, so we need to copy to a new buffer before
        # serializing; this slice of the underlying buffer (not the array) will ensure
        # that the buffer is properly copied at pickle-time.
        buf = buf.slice(byte_offset, byte_length)
    return buf


def _copy_bitpacked_buffer_if_needed(
    buf: "pyarrow.Buffer",
    offset: int,
    length: int,
) -> "pyarrow.Buffer":
    """Copy bit-packed binary buffer, if needed."""
    bit_offset = offset % 8
    byte_offset = offset // 8
    byte_length = _bytes_for_bits(bit_offset + length) // 8
    if offset > 0 or byte_length < buf.size:
        buf = buf.slice(byte_offset, byte_length)
        if bit_offset != 0:
            # Need to manually shift the buffer to eliminate the bit offset.
            buf = _align_bit_offset(buf, bit_offset, byte_length)
    return buf


def _copy_offsets_buffer_if_needed(
    buf: "pyarrow.Buffer",
    arr_type: "pyarrow.DataType",
    offset: int,
    length: int,
) -> Tuple["pyarrow.Buffer", int, int]:
    """Copy the provided offsets buffer, returning the copied buffer and the
    offset + length of the underlying data.
    """
    import pyarrow as pa
    import pyarrow.compute as pac

    if (
        pa.types.is_large_list(arr_type)
        or pa.types.is_large_string(arr_type)
        or pa.types.is_large_binary(arr_type)
        or pa.types.is_large_unicode(arr_type)
    ):
        offset_type = pa.int64()
    else:
        offset_type = pa.int32()
    # Copy offset buffer, if needed.
    buf = _copy_buffer_if_needed(buf, offset_type, offset, length + 1)
    # Reconstruct the offset array so we can determine the offset and length
    # of the child array.
    offsets = pa.Array.from_buffers(offset_type, length + 1, [None, buf])
    child_offset = offsets[0].as_py()
    child_length = offsets[-1].as_py() - child_offset
    # Create new offsets aligned to 0 for the copied data buffer slice.
    offsets = pac.subtract(offsets, child_offset)
    if pa.types.is_int32(offset_type):
        # We need to cast the resulting Int64Array back down to an Int32Array.
        offsets = offsets.cast(offset_type, safe=False)
    buf = offsets.buffers()[1]
    return buf, child_offset, child_length


def _bytes_for_bits(n: int) -> int:
    """Round up n to the nearest multiple of 8.
    This is used to get the byte-padded number of bits for n bits.
    """
    return (n + 7) & (-8)


def _align_bit_offset(
    buf: "pyarrow.Buffer",
    bit_offset: int,
    byte_length: int,
) -> "pyarrow.Buffer":
    """Align the bit offset into the buffer with the front of the buffer by shifting
    the buffer and eliminating the offset.
    """
    import pyarrow as pa

    bytes_ = buf.to_pybytes()
    bytes_as_int = int.from_bytes(bytes_, sys.byteorder)
    bytes_as_int >>= bit_offset
    bytes_ = bytes_as_int.to_bytes(byte_length, sys.byteorder)
    return pa.py_buffer(bytes_)


def _arrow_table_ipc_reduce(table: "pyarrow.Table"):
    """Custom reducer for Arrow Table that works around a zero-copy slicing pickling
    bug by using the Arrow IPC format for the underlying serialization.

    This is currently used as a fallback for unsupported types (or unknown bugs) for
    the manual buffer truncation workaround, e.g. for dense unions.
    """
    from pyarrow.ipc import RecordBatchStreamWriter
    from pyarrow.lib import BufferOutputStream

    output_stream = BufferOutputStream()
    with RecordBatchStreamWriter(output_stream, schema=table.schema) as wr:
        wr.write_table(table)
    # NOTE: output_stream.getvalue() materializes the serialized table to a single
    # contiguous bytestring, resulting in a few copy. This adds 1-2 extra copies on the
    # serialization side, and 1 extra copy on the deserialization side.
    return _restore_table_from_ipc, (output_stream.getvalue(),)


def _restore_table_from_ipc(buf: bytes) -> "pyarrow.Table":
    """Restore an Arrow Table serialized to Arrow IPC format."""
    from pyarrow.ipc import RecordBatchStreamReader

    with RecordBatchStreamReader(buf) as reader:
        return reader.read_all()


def _is_dense_union(type_: "pyarrow.DataType") -> bool:
    """Whether the provided Arrow type is a dense union."""
    import pyarrow as pa

    return pa.types.is_union(type_) and type_.mode == "dense"