Repository: huggingface/datasets
File: src/datasets/table.py (Python)

Code before:
        if isinstance(array_type, pa.PyExtensionType):
            return array_type.wrap_array(_concat_arrays([array.storage for array in arrays]))
        elif pa.types.is_struct(array_type):
            return pa.StructArray.from_arrays(
                [_concat_arrays([array.field(field.name) for array in arrays]) for field in array_type],
                fields=list(array_type),
                mask=pa.concat_arrays([array.is_null() for array in arrays]),
            )
        elif pa.types.is_list(array_type):
            if any(array.null_count > 0 for array in arrays):
                if config.PYARROW_VERSION.major < 10:
                    warnings.warn(
                        "None values are converted to empty lists in `pyarrow<10.0.0` when concatenating list arrays with None values. Install `pyarrow>=10.0.0` to avoid this behavior. More info: https://github.com/huggingface/datasets/issues/3676."
                    )
                else:
                    return pa.ListArray.from_arrays(
                        _offsets_concat([array.offsets for array in arrays]),
                        _concat_arrays([array.values for array in arrays]),
                        mask=pa.concat_arrays([array.is_null() for array in arrays]),
                    )
            return pa.ListArray.from_arrays(
                _offsets_concat([array.offsets for array in arrays]),
                _concat_arrays([array.values for array in arrays]),
            )
        elif pa.types.is_fixed_size_list(array_type):
            if config.PYARROW_VERSION.major < 13:
                # PyArrow bug: https://github.com/apache/arrow/issues/35360
                return pa.FixedSizeListArray.from_arrays(
                    _concat_arrays([array.values[array.offset * array.type.list_size :] for array in arrays]),
                    array_type.list_size,
                )
            else:
                return pa.FixedSizeListArray.from_arrays(
                    _concat_arrays([array.values for array in arrays]),
                    array_type.value_type,
                    array_type.list_size,
                )
        return pa.concat_arrays(arrays)

    return _concat_arrays(arrays)


@_wrap_for_chunked_arrays
def array_cast(array: pa.Array, pa_type: pa.DataType, allow_number_to_str=True):
    """Improved version of `pa.Array.cast`

    It supports casting `pa.StructArray` objects to re-order the fields.
    It also let you control certain aspects of the casting, e.g. whether
    to disable numbers (`floats` or `ints`) to strings.

    Args:

Review comment (bug):
Should we have something smarter ? like a function that tests if the pyarrow bug is still there and caches the result

Code after:
        if isinstance(array_type, pa.PyExtensionType):
            return array_type.wrap_array(_concat_arrays([array.storage for array in arrays]))
        elif pa.types.is_struct(array_type):
            return pa.StructArray.from_arrays(
                [_concat_arrays([array.field(field.name) for array in arrays]) for field in array_type],
                fields=list(array_type),
                mask=pa.concat_arrays([array.is_null() for array in arrays]),
            )
        elif pa.types.is_list(array_type):
            if any(array.null_count > 0 for array in arrays):
                if config.PYARROW_VERSION.major < 10:
                    warnings.warn(
                        "None values are converted to empty lists in `pyarrow<10.0.0` when concatenating list arrays with None values. Install `pyarrow>=10.0.0` to avoid this behavior. More info: https://github.com/huggingface/datasets/issues/3676."
                    )
                else:
                    return pa.ListArray.from_arrays(
                        _offsets_concat([array.offsets for array in arrays]),
                        _concat_arrays([array.values for array in arrays]),
                        mask=pa.concat_arrays([array.is_null() for array in arrays]),
                    )
            return pa.ListArray.from_arrays(
                _offsets_concat([array.offsets for array in arrays]),
                _concat_arrays([array.values for array in arrays]),
            )
        elif pa.types.is_fixed_size_list(array_type):
            if config.PYARROW_VERSION.major < 14:
                # PyArrow bug: https://github.com/apache/arrow/issues/35360
                return pa.FixedSizeListArray.from_arrays(
                    _concat_arrays([array.values[array.offset * array.type.list_size :] for array in arrays]),
                    array_type.list_size,
                )
            else:
                return pa.FixedSizeListArray.from_arrays(
                    _concat_arrays([array.values for array in arrays]),
                    array_type.value_type,
                    array_type.list_size,
                )
        return pa.concat_arrays(arrays)

    return _concat_arrays(arrays)


@_wrap_for_chunked_arrays
def array_cast(array: pa.Array, pa_type: pa.DataType, allow_number_to_str=True):
    """Improved version of `pa.Array.cast`

    It supports casting `pa.StructArray` objects to re-order the fields.
    It also let you control certain aspects of the casting, e.g. whether
    to disable numbers (`floats` or `ints`) to strings.

    Args: