anagrams kata (pyspark)

anagrams.py

from pyspark.sql import SparkSession
from pyspark.sql.types import StringType

from pyspark.sql.functions import col, udf, collect_list, size, length

spark = SparkSession.builder.getOrCreate()


def alpha_sort(word):
    return "".join(sorted(word))


sortedUdf = udf(alpha_sort, StringType())


df = (
    spark.read.csv("wordlist.txt")
    .withColumnRenamed("_c0", "word")
    .withColumn("sorted_word", sortedUdf(col("word")))
)

anagram_df = (
    df.groupBy("sorted_word")
    .agg(collect_list("word").alias("anagrams"))
    .where(size(col("anagrams")) > 1)
)

print(anagram_df.count())

print(anagram_df.orderBy(size(col("anagrams")).desc()).first())

print(anagram_df.orderBy(length(col("sorted_word")).desc()).first(10))