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

Code before:
        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)

    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)
                    and column.type != unified_struct_type
                ):
                    # 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]


Review comment (suggestion):
block_schema_field_names
```suggestion
        block_schema_field_names = set(block.schema.names)
```

Code after:
        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)

    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_schema_field_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_schema_field_names:
                column = block[column_name]

                # Check if the column type matches a struct type
                if (
                    isinstance(column.type, pa.StructType)
                    and column.type != unified_struct_type
                ):
                    # 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_schema_field_names]