from pyspark.sql import SparkSession
from pyspark.sql.functions import col
from pyspark.sql.dataframe import DataFrame

def main() -> None:
    print("Start of main function")

    # Create a Spark session using Spark Connect
    spark: SparkSession = SparkSession.builder \
        .appName("SparkConnectExample") \
        .remote("sc://localhost:50051") \
        .getOrCreate()
    print("Spark session created")

    # Create a DataFrame with a range of numbers
    # spark.range(5) creates a DataFrame with a single column named "id"
    # containing values from 0 to 4 (5 rows in total)
    df: DataFrame = spark.range(5)
    print("Created DataFrame with range(5)")
    # The DataFrame now looks like this:
    # +---+
    # | id|
    # +---+
    # |  0|
    # |  1|
    # |  2|
    # |  3|
    # |  4|
    # +---+

    # Write the DataFrame to a Parquet file
    df.write.mode("overwrite").parquet("myparquet.parquet")
    print("DataFrame written to Parquet file: myparquet.parquet")
    # The Parquet file contains the same data as the DataFrame above

    # Read the Parquet file back into a DataFrame
    df: DataFrame = spark.read.parquet("myparquet.parquet")
    print("DataFrame read from Parquet file")
    # The DataFrame remains unchanged:
    # +---+
    # | id|
    # +---+
    # |  0|
    # |  1|
    # |  2|
    # |  3|
    # |  4|
    # +---+

    print("DataFrame schema:")
    df.printSchema()
    # root
    #  |-- id: long (nullable = false)

    print("\nDataFrame content:")
    df.show()

    # Perform operations on the DataFrame
    # 1. filter(col("id") > 2): Select only rows where 'id' is greater than 2
    # 2. withColumn("id2", col("id") + 2): Add a new column 'id2' that is 'id' plus 2
    result: DataFrame = df.filter(col("id") > 2).withColumn("id2", col("id") + 2)

    print("\nFiltered and transformed DataFrame:")
    result.show()
    # The resulting DataFrame looks like this:
    # +---+---+
    # | id|id2|
    # +---+---+
    # |  3|  5|
    # |  4|  6|
    # +---+---+
    # Explanation:
    # 1. Only rows with id > 2 are kept (3 and 4)
    # 2. A new column 'id2' is added with values id + 2

    # Stop the Spark session
    spark.stop()
    print("Spark session stopped")

if __name__ == "__main__":
    main()