PySpark Reference

from pyspark.sql import SparkSession

# Initialize a SparkSession
spark = SparkSession.builder \
    .appName("PySparkReference") \
    .config("spark.some.config.option", "some-value") \
    .getOrCreate()

# Access the underlying SparkContext
sc = spark.sparkContext

# Get and set runtime configuration
spark.conf.set("spark.sql.shuffle.partitions", "10")
current_val = spark.conf.get("spark.sql.shuffle.partitions")
print(f"Partitions: {current_val}")

df = spark.createDataFrame([(1, "foo"), (2, "bar")], ["id", "val"])

# Display the first 20 rows (default)
df.show()

# Display with vertical formatting and no truncation
df.show(n=5, truncate=False, vertical=True)

# Read CSV with header and schema inference
df_csv = spark.read.csv("data.csv", header=True, inferSchema=True)

# Read Parquet (columnar format)
df_parquet = spark.read.parquet("data.parquet")

# Write to Parquet with overwrite mode
df.write.mode("overwrite").parquet("output.parquet")

# Write to CSV with header and custom delimiter
df.write.option("header", "true").option("sep", "|").csv("output.csv")

# Create DataFrame from a list of tuples or dictionaries
data = [("Alice", 34), ("Bob", 45)]
df = spark.createDataFrame(data, ["name", "age"])

from pyspark.sql import types as T

# Define explicit schema using StructType
schema = T.StructType([
    T.StructField("name", T.StringType(), True),
    T.StructField("age", T.IntegerType(), True),
    T.StructField("metadata", T.MapType(T.StringType(), T.StringType()), True)
])

df_with_schema = spark.createDataFrame([], schema)

from pyspark.sql import functions as F
from pyspark.sql import types as T

# Cast an existing column to a different type
df_cast = df.withColumn("age", F.col("age").cast(T.DoubleType()))

# Casting via selectExpr (SQL style)
df_cast_sql = df.selectExpr("cast(age as string) as age_str")

# Print schema in a hierarchical tree format
df.printSchema()

# Get column names and their data types as a list of tuples
print(df.dtypes)

from pyspark.sql import functions as F

# Select columns and rename using alias
df_sel = df.select(
    F.col("name").alias("user_name"),
    F.col("age") + 1
)

from pyspark.sql import functions as F

# Filter rows based on conditions (filter and where are synonyms)
df_filtered = df.filter(F.col("age") > 30)
df_where = df.where("age > 30 AND name LIKE 'A%'")

# Remove columns from DataFrame
df_less = df.drop("age")

# Remove duplicate rows based on all or specific columns
df_unique = df.dropDuplicates(["name"])

from pyspark.sql import functions as F

# Add a new column or replace an existing one
df_new = df.withColumn("is_adult", F.col("age") >= 18)

from pyspark.sql import functions as F

# Conditional logic similar to IF-THEN-ELSE
df_logic = df.withColumn("category", 
    F.when(F.col("age") < 18, "Minor")
     .when(F.col("age") < 65, "Adult")
     .otherwise("Senior")
)

from pyspark.sql import functions as F

# String and Date manipulations
df_ops = df.withColumn("upper_name", F.upper(F.col("name"))) \
           .withColumn("plus_7_days", F.date_add(F.current_date(), 7)) \
           .withColumn("diff", F.datediff(F.current_date(), F.to_date(F.lit("2020-01-01"))))

# Group DataFrame by one or more columns
grouped_df = df.groupBy("name")

from pyspark.sql import functions as F

# Perform aggregations on groups
df_agg = df.groupBy("name").agg(
    F.count("*").alias("total"),
    F.avg("age").alias("avg_age"),
    F.collect_list("age").alias("all_ages")
)

# Pivot column values into separate columns
df_pivot = df.groupBy("name").pivot("age").count()

df1 = spark.createDataFrame([(1, "A"), (2, "B")], ["id", "val1"])
df2 = spark.createDataFrame([(1, "X"), (3, "Y")], ["id", "val2"])

# Join types: inner, left, right, outer, left_semi, left_anti
df_inner = df1.join(df2, on="id", how="inner")
df_left = df1.join(df2, on="id", how="left")

from pyspark.sql import functions as F

# Force broadcasting a small DataFrame for performance
df_joined = df1.join(F.broadcast(df2), "id")

from pyspark.sql import functions as F

# Combine DataFrames vertically (must have same schema)
df_union = df1.union(df2.select("id", F.col("val2").alias("val1")))

# Find common rows
df_intersect = df1.intersect(df2.select("id", F.col("val2").alias("val1")))

from pyspark.sql.window import Window
from pyspark.sql import functions as F

# Define a window specification
window_spec = Window.partitionBy("category").orderBy(F.col("salary").desc())

from pyspark.sql import functions as F

# Apply ranking functions over the window
df_ranked = df.withColumn("rank", F.rank().over(window_spec)) \
              .withColumn("row_num", F.row_number().over(window_spec))

from pyspark.sql import functions as F

# Access values from preceding or following rows
df_trends = df.withColumn("prev_val", F.lag("salary", 1).over(window_spec)) \
              .withColumn("next_val", F.lead("salary", 1).over(window_spec))

from pyspark.sql import functions as F
from pyspark.sql import types as T

# Define and register a Python function as a UDF
@F.udf(returnType=T.StringType())
def complex_logic(val):
    return f"Processed: {val}"

df_udf = df.withColumn("result", complex_logic(F.col("name")))

import pandas as pd
from pyspark.sql import functions as F
from pyspark.sql import types as T

# Pandas UDF for vectorized execution (requires pyarrow)
@F.pandas_udf(T.DoubleType())
def vectorized_add_one(s: pd.Series) -> pd.Series:
    return s + 1

df_vec = df.withColumn("age_plus_one", vectorized_add_one(F.col("age")))

# Convert DataFrame to RDD
rdd = df.rdd

# Perform RDD operations
rdd_mapped = rdd.map(lambda row: (row.name, row.age * 2))

# Convert RDD back to DataFrame
df_back = rdd_mapped.toDF(["name", "double_age"])

from pyspark import StorageLevel

# Cache in memory (short-cut for persist(MEMORY_AND_DISK))
df.cache()

# Persist with explicit storage level
df.persist(StorageLevel.DISK_ONLY)

# Remove from cache
df.unpersist()

# Increase or shuffle partitions
df_repartitioned = df.repartition(100, "name")

# Decrease partitions without full shuffle
df_coalesced = df.coalesce(1)

# Display physical execution plan
df.explain()

# Display detailed plans (Parsed, Analyzed, Optimized, Physical)
df.explain(extended=True)