The Definitive Data Engineering Guide to Polars

from pyspark.sql.types import StructType, StructField, StringType, IntegerType

# Verbose and object-heavy schema definition
schema = StructType([
    StructField("user_id", IntegerType(), True),
    StructField("status", StringType(), True)
])

df = spark.read.schema(schema).parquet("s3://lake/events/")

import polars as pl

# Clean, dictionary-based schema definition
schema = {
    "user_id": pl.Int32,
    "status": pl.Categorical  # Highly optimized for low-cardinality strings
}

# scan_parquet creates a LazyFrame. No data is read into RAM yet.
lazy_df = pl.scan_parquet("s3://lake/events/*.parquet", schema=schema)

import pyspark.sql.functions as F

df_clean = df.withColumn(
    "email_clean", 
    F.trim(F.lower(F.col("email")))
).withColumn(
    "status",
    F.when(F.col("age") < 18, "Minor")
     .otherwise("Adult")
).filter(F.col("status") == "Adult")

import polars as pl

df_clean = df.with_columns(
    # Chained method syntax is cleaner than nested functions
    pl.col("email").str.to_lowercase().str.strip_chars().alias("email_clean"),
    
    pl.when(pl.col("age") < 18)
      .then(pl.lit("Minor"))
      .otherwise(pl.lit("Adult"))
      .alias("status")
).filter(pl.col("status") == "Adult")

metrics_df = df.groupBy("region", "department").agg(
    F.sum("revenue").alias("total_revenue"),
    F.countDistinct("user_id").alias("unique_users"),
    F.first("event_time").alias("first_event")
)

metrics_df = df.group_by("region", "department").agg(
    pl.col("revenue").sum().alias("total_revenue"),
    pl.col("user_id").n_unique().alias("unique_users"),
    pl.col("event_time").first().alias("first_event")
)

from pyspark.sql.window import Window

# Define the window specification
w_rank = Window.partitionBy("department").orderBy(F.desc("salary"))
w_avg = Window.partitionBy("department")

ranked_df = df.withColumn(
    "dept_rank", 
    F.rank().over(w_rank)
).withColumn(
    "dept_avg_salary",
    F.mean("salary").over(w_avg)
)

ranked_df = df.with_columns(
    # Rank users by salary within their department
    pl.col("salary")
      .rank(descending=True)
      .over("department")
      .alias("dept_rank"),
      
    # Attach the department average to every row
    pl.col("salary")
      .mean()
      .over("department")
      .alias("dept_avg_salary")
)

# Doing this in Spark safely usually requires dropping to SQL
# or using an inefficient complex join condition.

trades.createOrReplaceTempView("trades")
quotes.createOrReplaceTempView("quotes")

joined_df = spark.sql("""
    SELECT t.*, q.price AS quote_price
    FROM trades t
    LEFT JOIN quotes q ON t.ticker = q.ticker
      AND q.timestamp <= t.timestamp
    -- Requires subsequent windowing to keep only the latest row!
    -- VERY SLOW.
""")

# DataFrames must be sorted by the 'on' key first.
# Executes an ultra-fast O(N) backward search.

joined_df = trades_df.join_asof(
    quotes_df,
    on="timestamp",       # The temporal matching column
    by="ticker",          # Exact match grouping condition
    strategy="backward",  # Find closest quote PRIOR to trade
    tolerance="5m"        # Reject matches older than 5 mins
)

# Finding if 'admin' is in the tags array
df = df.withColumn(
    "is_admin",
    F.array_contains(F.col("tags"), "admin")
)

# Exploding an array into multiple rows
exploded_df = df.select(
    "user_id", 
    F.explode("tags").alias("tag")
)

# The native `.list` namespace provides array operations
df = df.with_columns(
    pl.col("tags").list.contains("admin").alias("is_admin")
)

# Exploding is a native DataFrame method
exploded_df = df.explode("tags")