Reshuffle Schema Columns in Spark DataFrame

 1def move_col_after(df: DataFrame, col_to_move: str, col_after: str) -> DataFrame:
 2    """Moves a column to appear after specified column in a dataframe
 3
 4    :param df: Dataframe to work with
 5    :param col_to_move: column that is to be moved
 6    :param col_after: column to move after
 7    :return:
 8    """
 9    ordered_cols = []
10    moved = False
11    for fld in df.schema.fieldNames():
12
13        if fld != col_to_move:
14            ordered_cols.append(fld)
15
16        if fld == col_after:
17            ordered_cols.append(col_to_move)
18            moved = True
19
20    # in case the column wasn't found, don't lose the data
21    if not moved:
22        ordered_cols.append(col_to_move)
23
24    return df.select(*ordered_cols)
25
26
27def move_cols_after(df: DataFrame, col_after: str, *cols_to_move) -> DataFrame:
28    """Moves a column to appear after specified column in a dataframe
29
30    :param df: Dataframe to work with
31    :param col_after: column to move after
32    :param cols_to_move: columns to be moved
33    :return:
34    """
35
36    if not cols_to_move:
37        return df
38
39    ordered_cols = []
40    for fld in df.schema.fieldNames():
41
42        if fld not in cols_to_move:
43            ordered_cols.append(fld)
44
45        if fld == col_after:
46            ordered_cols.extend(cols_to_move)
47    return df.select(*ordered_cols)
48
49
50def move_col_to_position(df: DataFrame, col_name: str, pos: int) -> DataFrame:
51    rc = []
52    for i in range(0, len(df.columns)):
53        if i == pos:
54            rc.append(col_name)
55
56        col = df.columns[i]
57        if col != col_name:
58            rc.append(col)
59
60    return df.select(*rc)