Get feature permutations, one for each row, using pyspark

pyspark_get_feature_permutations.py

import pyspark
from pyspark.sql import functions as F


def get_features_permutations(
        df: pyspark.DataFrame,
        feature_names: list,
        output_col='features_permutations'
):
    """
    Creates a column for the ordered features and then shuffles it.
    The result is a dataframe with a column `output_col` that contains:
    [feat2, feat4, feat3, feat1],
    [feat3, feat4, feat2, feat1],
    [feat1, feat2, feat4, feat3],
    ...
    """
    return df.withColumn(
        output_col,
        F.shuffle(
            F.array(*[F.lit(f) for f in feature_names])
        )
    )