Repository: ray-project/ray
File: python/ray/data/_internal/arrow_ops/transform_pyarrow.py (Python)

Code before:
    except Exception as e:
        raise ArrowConversionError(
            f"Failed to unify schemas: {str(e)}\n"
            f"{'-' * 16}\n"
            f"Schemas:\n"
            f"{'-' * 16}\n"
            f"{schemas_to_unify}"
        ) from e

    # Handle alignment of struct type columns.
    blocks = _align_struct_fields(blocks, schema)

    # Identify columns with null lists
    cols_with_null_list = set()
    for b in blocks:
        for col_name in b.schema.names:
            col_type = b.schema.field(col_name).type
            if pa.types.is_list(col_type) and pa.types.is_null(col_type.value_type):
                cols_with_null_list.add(col_name)

    # Concatenate the columns according to their type
    concatenated_cols = {}
    native_pyarrow_cols = []
    for col_name in schema.names:
        col_type = schema.field(col_name).type

        # _align_struct_fields guarantees the existence of the column
        col_chunked_arrays = [block.column(col_name) for block in blocks]

        if col_name in cols_with_null_list:
            concatenated_cols[col_name] = _concat_cols_with_null_list(
                col_chunked_arrays
            )
        elif isinstance(col_type, tensor_types):
            concatenated_cols[col_name] = _concat_cols_with_extension_tensor_types(
                col_chunked_arrays
            )
        elif isinstance(col_type, ArrowPythonObjectType):
            concatenated_cols[col_name] = _concat_cols_with_extension_object_types(
                col_chunked_arrays
            )
        else:
            # Add to the list of native pyarrow columns, these will be concatenated after the loop using pyarrow.concat_tables
            native_pyarrow_cols.append(col_name)

    concatenated_cols.update(
        _concat_cols_with_native_pyarrow_types(
            native_pyarrow_cols, blocks, promote_types
        )
    )


Review comment (bug):
This actually fixes another subtle bug as well.

Previously, when a block did not contain `col_name`, we were appending
`pa.nulls(block.num_rows, type=col_type)` to `col_chunked_arrays`. That returns
a plain `Array` instead of a `ChunkedArray`, which can cause
`_concat_cols_with_null_list` to raise when it tries to access `arr.chunks`.

Code after:
            f"Failed to unify schemas: {str(e)}\n"
            f"{'-' * 16}\n"
            f"Schemas:\n"
            f"{'-' * 16}\n"
            f"{schemas_to_unify}"
        ) from e

    # Handle alignment of struct type columns.
    blocks = _align_struct_fields(blocks, schema)

    # Identify columns with null lists
    cols_with_null_list = set()
    for b in blocks:
        for col_name in b.schema.names:
            col_type = b.schema.field(col_name).type
            if pa.types.is_list(col_type) and pa.types.is_null(col_type.value_type):
                cols_with_null_list.add(col_name)

    # Concatenate the columns according to their type
    concatenated_cols = {}
    native_pyarrow_cols = []
    for col_name in schema.names:
        col_type = schema.field(col_name).type

        # _align_struct_fields guarantees the existence of the column
        col_chunked_arrays = [block.column(col_name) for block in blocks]

        if col_name in cols_with_null_list:
            concatenated_cols[col_name] = _concat_cols_with_null_list(
                col_chunked_arrays
            )
        elif isinstance(col_type, tensor_types):
            concatenated_cols[col_name] = _concat_cols_with_extension_tensor_types(
                col_chunked_arrays
            )
        elif isinstance(col_type, ArrowPythonObjectType):
            concatenated_cols[col_name] = _concat_cols_with_extension_object_types(
                col_chunked_arrays
            )
        else:
            # Add to the list of native pyarrow columns, these will be concatenated after the loop using pyarrow.concat_tables
            native_pyarrow_cols.append(col_name)

    concatenated_cols.update(
        _concat_cols_with_native_pyarrow_types(
            native_pyarrow_cols, blocks, promote_types
        )
    )

    # Ensure that the columns are in the same order as the schema, reconstruct the table.
    return pyarrow.Table.from_arrays(