PyFlow Analysis Report

from pyspark.sql import SparkSession

from pyspark.sql.functions import col,

lit, when, udf

from pyspark.sql.types import IntegerType, StringType

spark = SparkSession.builder().appName("LineageTest").getOrCreate()

data1 = [("Alice", 34), ("Bob", 45), ("Charlie", 23)]

columns = ["name", "age"]

df1 = spark.createDataFrame(data1, columns)

data2 = [("Dave", 29), ("Eve", 52)]

df2 = spark.createDataFrame(data2, columns)

def age_category(age):

if age < 30:
        return "Young"
    else:
        return "Old"

age_category_udf = udf(age_category, StringType())

def add_age_category(input_df):
    return (
            input_df
            .withColumn("age_category", age_category_udf(col("age")))
        )

def filter_and_select(input_df, min_age):
    return (
            input_df
            .filter(col("age") > min_age)
            .select(
                col("name"),
                col("age_category"),
                (col("age") * lit(2)).alias("double_age")
            )
        )

def aggregate_stats(input_df):
    return (
            input_df
            .groupBy("age_category")
            .agg(
                {"age": "avg", "*": "count"}
            )
            .withColumnRenamed("avg(age)", "avg_age")
            .withColumnRenamed("count(1)", "count")
        )

def repartition_and_cache(input_df, num_partitions):
    return input_df.repartition(num_partitions).cache()

if __name__ == "__main__":

df1_cat = add_age_category(df1)

df2_cat = add_age_category(df2)

df_union = (
            df1_cat.union(df2_cat)
            .dropDuplicates(["name"])
            .orderBy("name")
        )

df_filtered = filter_and_select(df_union, min_age=30)

df_stats = aggregate_stats(df1_cat)

df_stats_cached = repartition_and_cache(df_stats, num_partitions=2)

df_joined = (
            df_filtered
            .join(df_stats_cached, on="age_category", how="left")
            .orderBy(col("count").desc())
            .limit(10)
        )

df_joined.show(truncate=False)

spark.stop()