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

Code before:
            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 schema matches exactly, skip
        if block.schema.equals(schema):
            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
            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

        col_chunked_arrays = []
        for block in blocks:
            if col_name in block.schema.names:
                col_chunked_arrays.append(block.column(col_name))
            else:
                col_chunked_arrays.append(pa.nulls(block.num_rows, type=col_type))


Review comment (performance):
Actually I realized one more thing while reviewing. You can actually just incorporate this into `_align_struct_fields`.

```
def _align_struct_fields(
    blocks: List["pyarrow.Table"], schema: "pyarrow.Schema"
) -> List["pyarrow.Table"]:
    """
    Align struct columns across blocks to match the provided schema.

    Args:
        blocks: List of Arrow tables to align.
        schema: Unified schema with desired struct column alignment.

    Returns:
        List[pa.Table]: List of aligned Arrow tables.
    """
    import pyarrow as pa

    # Check if all block schemas are already aligned
    if all(block.schema == schema for block in blocks):
        return blocks

    # Extract all struct column types from the provided schema
    unified_struct_types = _extract_unified_struct_types(schema)

    # If there are no struct columns in the schema, return blocks as is
    if not unified_struct_types:
        return blocks

    aligned_blocks = []

    # Iterate over each block (table) in the list
    for block in blocks:
        # Fast path: if schema matches exactly, skip
        if block.schema.equals(schema):
            aligned_blocks.append(block)
            continue

        # Get the number of rows in the block
        block_length = len(block)
        block_names = set(block.schema.names)

        # Process struct columns that need alignment
        for column_name, unified_struct_type in unified_struct_types.items():
            if column_name in block_names:
                column = block[column_name]
                # Check if the column type matches a struct type
                if isinstance(column.type, pa.StructType):
                    # Align struct fields
                    aligned_column = _backfill_missing_fields(
                        column, unified_struct_type, block_length
                    )
                    # Replace the column with the aligned version
                    block = block.set_column(
                        block.schema.get_field_index(column_name),
                        column_name,
                        aligned_column,
                    )

        # Find all missing columns (struct and non-struct)
        missing_fields = [f for f in schema if f.name not in block_names]

        # Append missing columns with null values
        # 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(block_length, type=field.type)
                block = block.append_column(field.name, null_col)

        # Reorder columns to match the target schema
        # select() is a zero-copy operation for column reordering
        block = block.select(schema.names)
        aligned_blocks.append(block)

    # Return the list of aligned blocks
    return aligned_blocks

```



Code after:
            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(
        [concatenated_cols[col_name] for col_name in schema.names], schema=schema