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

Code before:
    if not blocks:
        # Short-circuit on empty list of blocks.
        return pa.table([])

    if len(blocks) == 1:
        return blocks[0]

    # If the result contains pyarrow schemas, unify them
    schemas_to_unify = [b.schema for b in blocks]
    try:
        schema = unify_schemas(schemas_to_unify, promote_types=promote_types)
    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

    # Backfill missing columns
    blocks = list(blocks)
    for i, block in enumerate(blocks):
        # Fast path: if there is no missing column and column order matches, skip
        # Struct field alignment is handled separately in _align_struct_fields.
        if block.schema.names == schema.names:
            continue

        # 1. Find fields that are in the target schema but missing from the block
        block_names = set(block.schema.names)
        missing_fields = [f for f in schema if f.name not in block_names]

        # 2. Create and append null columns for missing fields
        # This avoids rebuilding the whole table and only adds what's needed
        if missing_fields:
            for field in missing_fields:
                # pa.nulls creates a null array of the correct length and type efficiently
                null_col = pa.nulls(len(block), type=field.type)
                block = block.append_column(field.name, null_col)

        # 3. Reorder columns to match the target schema
        blocks[i] = block.select(schema.names)

    # 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

Review comment (performance):
This is unfortunately much slower because it does element wise comparison in python. The `.equals()` method for the pyarrow schema object does the checking in C++ 

Code after:

    if not blocks:
        # Short-circuit on empty list of blocks.
        return pa.table([])

    if len(blocks) == 1:
        return blocks[0]

    # If the result contains pyarrow schemas, unify them
    schemas_to_unify = [b.schema for b in blocks]
    try:
        schema = unify_schemas(schemas_to_unify, promote_types=promote_types)
    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(