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

Code before:
                if isinstance(column.type, pa.StructType):
                    aligned_columns[column_name] = _backfill_missing_fields(
                        column, unified_struct_type, block_length
                    )
                else:
                    # If the column is not a struct, simply keep the original column
                    aligned_columns[column_name] = column
            else:
                # If the column is missing, create a null-filled column with the same
                # length as the block
                aligned_columns[column_name] = pa.array(
                    [None] * block_length, type=unified_struct_type
                )

        # Create a new aligned block with the updated columns and the unified schema.
        new_columns = []
        for column_name in schema.names:
            if column_name in aligned_columns:
                # Use the aligned column if available
                new_columns.append(aligned_columns[column_name])
            elif column_name in block.schema.names:
                # Use the original column if not aligned
                new_columns.append(block[column_name])
            else:
                null_array = pa.nulls(len(block), type=schema.field(column_name).type)
                new_columns.append(null_array)
        aligned_blocks.append(pa.table(new_columns, schema=schema))

    # Return the list of aligned blocks
    return aligned_blocks


def shuffle(block: "pyarrow.Table", seed: Optional[int] = None) -> "pyarrow.Table":
    """Shuffles provided Arrow table"""

    if len(block) == 0:
        return block

    indices = np.arange(block.num_rows)
    # Shuffle indices
    np.random.RandomState(seed).shuffle(indices)

    return take_table(block, indices)


def _concat_cols_with_null_list(
    col_chunked_arrays: List["pyarrow.ChunkedArray"],
) -> "pyarrow.ChunkedArray":
    import pyarrow as pa

    # For each opaque list column, iterate through all schemas until

Review comment (refactor):
Can you please use `BlockAccessor` abstraction for this operation?

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

        # Append missing columns with null values
        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


def shuffle(block: "pyarrow.Table", seed: Optional[int] = None) -> "pyarrow.Table":
    """Shuffles provided Arrow table"""

    if len(block) == 0:
        return block

    indices = np.arange(block.num_rows)
    # Shuffle indices
    np.random.RandomState(seed).shuffle(indices)

    return take_table(block, indices)


def _concat_cols_with_null_list(
    col_chunked_arrays: List["pyarrow.ChunkedArray"],
) -> "pyarrow.ChunkedArray":
    import pyarrow as pa

    # For each opaque list column, iterate through all schemas until
    # we find a valid value_type that can be used to override the
    # column types in the following for-loop.