pyspark code

case

find adult or not

from pyspark.sql.functions import when, col

df = df.withColumn(
    "age_group",
    when(col("age").isNull(),  "unknown")
    .when(col("age") < 0,      "invalid")
    .when(col("age") < 18,     "minor")
    .otherwise("adult")
)

filter

# by condition
df.filter(col("age") > 18)
df.filter("age > 18")                          # SQL string — same result

# multiple conditions
df.filter((col("age") > 18) & (col("city") == "pune"))
df.filter((col("age") > 18) | (col("city") == "pune"))
df.filter(~(col("city") == "pune"))            # NOT

# null checks
df.filter(col("age").isNull())
df.filter(col("age").isNotNull())

# string
df.filter(col("name").startswith("A"))
df.filter(col("name").endswith("e"))
df.filter(col("name").contains("ali"))
df.filter(col("name").like("Ali%"))            # SQL LIKE, % wildcard
df.filter(col("name").rlike("^[Aa]li.*"))      # regex

# in / not in
df.filter(col("city").isin("pune", "mumbai"))
df.filter(~col("city").isin("pune", "mumbai"))

# between (inclusive)
df.filter(col("age").between(18, 60))

# date
df.filter(col("dt") >= "2024-01-01")
df.filter(col("dt").between("2024-01-01", "2024-12-31"))

# on struct field
df.filter(col("address.city") == "pune")

# on array — check if value exists in array
from pyspark.sql.functions import array_contains
df.filter(array_contains(col("tags"), "premium"))

# after explode
df.filter(col("order.amt") > 500)

# SQL style
df.filter("city = 'pune' AND age > 18")
df.where("city = 'pune'")                      # where = filter, same thing

Gotchas

• & | ~ not and or not — Python keywords short-circuit, don't work on columns.
• Always wrap conditions in () when chaining — operator precedence bites without them.
• Filtering on a null with == returns false, not true — use isNull().
• String filter ("age > 18") is fine for simple cases; use col() when column name is dynamic.

window functions

from pyspark.sql.functions import (
    row_number, rank, dense_rank, ntile,
    lag, lead,
    sum, avg, min, max, count,
    first, last,
    percent_rank, cume_dist
)
from pyspark.sql.window import Window

# sample data: customer orders
# customer_id, order_date, amount, city

define windows

# partition + order
w = Window.partitionBy("customer_id").orderBy("order_date")

# partition only (for aggregations over whole group)
w_agg = Window.partitionBy("customer_id")

# running total frame
w_run = Window.partitionBy("customer_id").orderBy("order_date") \
              .rowsBetween(Window.unboundedPreceding, Window.currentRow)

# sliding window — last 3 rows including current
w_slide = Window.partitionBy("customer_id").orderBy("order_date") \
                .rowsBetween(-2, Window.currentRow)

# range frame — rows within 7 days of current row
w_range = Window.partitionBy("customer_id").orderBy("order_date_long") \
                .rangeBetween(-7 * 86400, Window.currentRow)   # epoch seconds

Ranking

python

df = df.withColumn("row_num",    row_number().over(w))   # no ties, always unique
df = df.withColumn("rank",       rank().over(w))         # ties get same rank, next rank skips
df = df.withColumn("dense_rank", dense_rank().over(w))   # ties get same rank, no skip
df = df.withColumn("ntile_4",    ntile(4).over(w))       # quartile bucket 1..4
df = df.withColumn("pct_rank",   percent_rank().over(w)) # 0.0 to 1.0
df = df.withColumn("cume_dist",  cume_dist().over(w))    # fraction of rows <= current

---

Lag / lead

python

df = df.withColumn("prev_amount",  lag("amount",  1).over(w))          # previous row
df = df.withColumn("next_amount",  lead("amount", 1).over(w))          # next row
df = df.withColumn("prev_2_amount",lag("amount",  2, 0).over(w))       # 2 rows back, default 0
df = df.withColumn("diff",         col("amount") - lag("amount", 1).over(w))  # change

---

Aggregations

python

# running
df = df.withColumn("running_total",  sum("amount").over(w_run))
df = df.withColumn("running_avg",    avg("amount").over(w_run))
df = df.withColumn("running_count",  count("amount").over(w_run))

# sliding (last 3 rows)
df = df.withColumn("sliding_avg", avg("amount").over(w_slide))

# whole partition
df = df.withColumn("total_per_customer",   sum("amount").over(w_agg))
df = df.withColumn("max_per_customer",     max("amount").over(w_agg))
df = df.withColumn("pct_of_customer_total",
                   col("amount") / sum("amount").over(w_agg))

---

First / last

python

# first/last need ignorNulls and explicit frame for reliable results
w_full = Window.partitionBy("customer_id").orderBy("order_date") \
               .rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing)

df = df.withColumn("first_order_amt", first("amount", ignorNulls=True).over(w_full))
df = df.withColumn("last_order_amt",  last("amount",  ignorNulls=True).over(w_full))

---

Common patterns

python

# latest record per customer (dedup)
from pyspark.sql.functions import col

deduped = (df.withColumn("rn", row_number().over(w))
             .filter("rn = 1")
             .drop("rn"))

# top N per group
top2 = (df.withColumn("rn", row_number().over(w))
          .filter(col("rn") <= 2)
          .drop("rn"))

# percent of total per partition
df = df.withColumn("pct", col("amount") / sum("amount").over(w_agg) * 100)

# flag if amount increased from previous row
df = df.withColumn("increased",
       col("amount") > lag("amount", 1).over(w))

# running total reset — gaps and islands
df = df.withColumn("grp", sum("is_new_session").over(w_run))

Frame cheatsheet

• rowsBetween - physical offset

• rangeBetween - value offset (needs numeric/date orderBy)

• unboundedPreceding

• unboundedFollowing
• currentRow

Gotchas

• No partitionBy → all data in one partition → OOM on large tables.
• Default frame with orderBy = RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW — includes ties.
• Default frame without orderBy = whole partition.
• last without explicit unboundedFollowing frame returns current row, not partition last.
• rank / dense_rank / row_number ignore frame — always rank within full partition order.
• Window functions not allowed in WHERE/HAVING — wrap in subquery or CTE.
• Multiple windows on same partition+order → Spark reuses one shuffle. Define them consistently.
• percent_rank of first row = 0.0 always; cume_dist of last row = 1.0 always.

Delta lake

read

# table name (Unity Catalog: catalog.schema.table)
df = spark.read.table("catalog.schema.orders")

# path
df = spark.read.format("delta").load("/mnt/delta/orders")

# time travel
df = spark.read.option("versionAsOf", 5).table("catalog.schema.orders")
df = spark.read.option("timestampAsOf", "2024-01-01").table("catalog.schema.orders")

write

# append
(
df
    .write
    .format("delta")
    .mode("append")
    .option("mergeSchema", "true") # schema evolution
    .saveAsTable("catalog.schema.orders")
)

# overwrite specific partitions only
(df.write.format("delta")
   .mode("overwrite")
   .option("replaceWhere", "dt >= '2024-01-01' AND dt < '2024-02-01'")
   .saveAsTable("catalog.schema.orders"))

# overwrite and replace schema entirely
(df.write.format("delta")
   .mode("overwrite")
   .option("overwriteSchema", "true")
   .saveAsTable("catalog.schema.orders"))

create table

# DDL
spark.sql("""
  CREATE TABLE IF NOT EXISTS catalog.schema.orders (
    order_id  BIGINT,
    customer  STRING,
    amt       DECIMAL(10,2),
    dt        DATE
  )
  USING DELTA
  CLUSTER BY (customer_id)
  TBLPROPERTIES (
    'delta.enableDeletionVectors' = 'true',
    'delta.enableChangeDataFeed'  = 'true'
  )
""")

maintenance

spark.sql("OPTIMIZE catalog.schema.orders")                        # compaction
spark.sql("OPTIMIZE catalog.schema.orders ZORDER BY (customer_id)")
spark.sql("VACUUM  catalog.schema.orders RETAIN 168 HOURS")        # 7 days
spark.sql("ANALYZE TABLE catalog.schema.orders COMPUTE STATISTICS FOR ALL COLUMNS")  # CBO stats
spark.sql("DESCRIBE HISTORY catalog.schema.orders")
spark.sql("DESCRIBE DETAIL  catalog.schema.orders")

kafka

read

raw = (spark.readStream
       .format("kafka")
       .option("kafka.bootstrap.servers", "broker1:9092,broker2:9092")
       .option("subscribe", "orders")
       .option("startingOffsets", "latest")        # latest / earliest / json offset
       .option("maxOffsetsPerTrigger", 100_000)    # pace ingestion
       .option("failOnDataLoss", "false")          # topic deleted / offsets expired
       .load())

batch read (backfill/one-time)

df = (spark.read                               # read not readStream
      .format("kafka")
      .option("kafka.bootstrap.servers", "broker1:9092,broker2:9092")
      .option("subscribe", "orders")
      .option("startingOffsets", """{"orders":{"0":1000,"1":1000}}""")  # exact offsets
      .option("endingOffsets",   """{"orders":{"0":2000,"1":2000}}""")
      .load())

columns you get:

Column	Type	Notes
key	binary	cast to STRING if text key
value	binary	your payload — always cast/parse
topic	string	useful when subscribing to multiple
partition	int
offset	long
timestamp	timestamp	producer or broker timestamp
timestampType	int	0=CreateTime, 1=LogAppendTime

parse value

from pyspark.sql.functions import col, from_json, from_avro
from pyspark.sql.types import StructType, StructField, StringType, IntegerType, TimestampType

schema = StructType([
    StructField("order_id", StringType()),
    StructField("customer", StringType()),
    StructField("amt",      IntegerType()),
    StructField("ts",       TimestampType())
])

# JSON value
parsed = (raw
          .select(from_json(col("value").cast("string"), schema).alias("d"),
                  col("timestamp").alias("kafka_ts"))
          .select("d.*", "kafka_ts"))

# Avro value (with schema registry)
parsed = (raw
          .select(from_avro(col("value"), "<avro_schema_json_string>").alias("d"))
          .select("d.*"))

write

streaming write

(parsed
 .selectExpr(
     "CAST(order_id AS STRING) AS key",   # key must be STRING or BINARY
     "to_json(struct(*)) AS value"         # value must be STRING or BINARY
 )
 .writeStream
 .format("kafka")
 .option("kafka.bootstrap.servers", "broker1:9092,broker2:9092")
 .option("topic", "orders_enriched")
 .option("checkpointLocation", "/mnt/checkpoints/orders_enriched")
 .outputMode("append")
 .trigger(processingTime="30 seconds")
 .start())

batch write

(df.selectExpr("CAST(id AS STRING) AS key", "to_json(struct(*)) AS value")
   .write
   .format("kafka")
   .option("kafka.bootstrap.servers", "broker1:9092,broker2:9092")
   .option("topic", "orders_out")
   .save())

csv

read

df = (spark.read
      .format("csv")
      .option("header", "true")
      .option("inferSchema", "true")       # avoid in prod — triggers extra scan
      .load("/mnt/data/orders.csv"))

# explicit schema (prod way)
schema = "order_id INT, customer STRING, amt DECIMAL(10,2), dt DATE"

df = (spark.read
      .format("csv")
      .schema(schema)
      .option("header", "true")
      .option("dateFormat", "yyyy-MM-dd")
      .option("timestampFormat", "yyyy-MM-dd HH:mm:ss")
      .option("nullValue", "NULL")         # treat this string as null
      .option("emptyValue", "")
      .option("mode", "PERMISSIVE")        # PERMISSIVE / DROPMALFORMED / FAILFAST
      .option("columnNameOfCorruptRecord", "_corrupt_record")
      .load("/mnt/data/orders.csv"))

# read folder — all CSVs in directory
df = spark.read.schema(schema).option("header","true").csv("/mnt/data/orders/")

# read multiple explicit paths
df = spark.read.schema(schema).option("header","true").csv(
    "/mnt/data/orders_jan.csv",
    "/mnt/data/orders_feb.csv"
)

options

Option	Default	Notes
header	false	use first row as column names
sep / delimiter	,	any single char; \t for TSV
quote	"	quoting character
escape	\	escape character inside quotes
multiLine	false	fields with newlines inside quotes
encoding	UTF-8	ISO-8859-1 for legacy files
ignoreLeadingWhiteSpace	false
ignoreTrailingWhiteSpace	false
nanValue	NaN	string to treat as NaN
positiveInf / negativeInf	Inf / -Inf
comment	disabled	skip lines starting with this char

write

# basic
(df.write
   .format("csv")
   .mode("overwrite")               # overwrite / append / error / ignore
   .option("header", "true")
   .save("/mnt/data/output/orders"))

# options
(df.write
   .format("csv")
   .mode("overwrite")
   .option("header", "true")
   .option("sep", ",")
   .option("quote", '"')
   .option("escape", "\\")
   .option("nullValue", "")
   .option("dateFormat", "yyyy-MM-dd")
   .option("timestampFormat", "yyyy-MM-dd HH:mm:ss")
   .option("compression", "none")   # none / gzip / bz2 / deflate / snappy
   .save("/mnt/data/output/orders"))

# control number of output files
df.coalesce(1).write.csv(...)       # single file — small data only
df.repartition(10).write.csv(...)   # 10 files

scd

scd 1

merge - upsert

from delta.tables import DeltaTable

target = DeltaTable.forName(spark, "catalog.schema.orders")

(target.alias("t")
       .merge(df.alias("s"), "t.order_id = s.order_id")
       .whenMatchedUpdateAll()
       .whenNotMatchedInsertAll()
       .whenNotMatchedBySourceDelete()   # deletes rows in target not in source
       .execute())

scd2

from pyspark.sql.functions import current_timestamp, lit
from delta.tables import DeltaTable

source = spark.read.table("catalog.schema.stg_customer")  # id, name, city


target = DeltaTable.forName(spark, "catalog.schema.dim_customer")

now = current_timestamp()

# step 1 — expire current rows where name or city changed
(target.alias("t")
       .merge(source.alias("s"), "t.id = s.id AND t.is_current = true")
       .whenMatchedUpdate(
           condition = "t.name <> s.name OR t.city <> s.city",
           set = {
               "valid_to":   "current_timestamp()",
               "is_current": "false"
           }
       )
       .execute())

# step 2 — insert new rows for changed + new records
(source.alias("s")
       .join(
           spark.read.table("catalog.schema.dim_customer")
                .filter("is_current = true")
                .alias("t"),
           on  = "s.id = t.id",
           how = "left"
       )
       .filter("t.id IS NULL OR t.name <> s.name OR t.city <> s.city")
       .select(
           "s.id", "s.name", "s.city",
           current_timestamp().alias("valid_from"),
           lit(None).cast("timestamp").alias("valid_to"),
           lit(True).alias("is_current")
       )
       .write.format("delta").mode("append")
       .saveAsTable("catalog.schema.dim_customer"))

single merge - dbr 12.2+

# works cleanly when Delta supports multi-action MERGE
target = DeltaTable.forName(spark, "catalog.schema.dim_customer")

staged = (source.alias("s")
                .join(
                    spark.read.table("catalog.schema.dim_customer")
                         .filter("is_current = true").alias("t"),
                    on  = "s.id = t.id",
                    how = "left"
                )
                .selectExpr(
                    "s.id",
                    "s.name",
                    "s.city",
                    "CASE WHEN t.id IS NULL OR t.name <> s.name OR t.city <> s.city THEN 'new' ELSE 'same' END AS action",
                    "t.id AS t_id"
                )
                .filter("action = 'new'"))

(target.alias("t")
       .merge(staged.alias("s"), "t.id = s.t_id AND t.is_current = true")
       .whenMatchedUpdate(set = {
           "valid_to":   "current_timestamp()",
           "is_current": "false"
       })
       .whenNotMatchedInsert(values = {
           "id":         "s.id",
           "name":       "s.name",
           "city":       "s.city",
           "valid_from": "current_timestamp()",
           "valid_to":   "null",
           "is_current": "true"
       })
       .execute())

Gotchas

• Two-step approach is safer — MERGE then INSERT avoids race conditions on concurrent writes.
• Never update valid_from on existing rows — only expire (valid_to + is_current).
• is_current = true filter on target before MERGE — avoids matching already-expired rows.
• Add surrogate_key (auto-increment or hash of id + valid_from) if fact tables need a stable FK.
• valid_to = null for current row is cleaner than a sentinel date (9999-12-31) — use IS NULL in queries.
• Partition by is_current or valid_to IS NULL helps if table is large and most queries hit current only.
• If source can have duplicate IDs, dedupe before MERGE: source.dropDuplicates(["id"]).

sqlserver

jdbc_url = "jdbc:sqlserver://server.database.windows.net:1433;databaseName=mydb"

connection_props = {
    "user":     dbutils.secrets.get("scope", "sql-user"),
    "password": dbutils.secrets.get("scope", "sql-password"),
    "driver":   "com.microsoft.sqlserver.jdbc.SQLServerDriver"
}

azure sql with aad token (no password)

import struct, pyodbc
from azure.identity import ManagedIdentityCredential

cred  = ManagedIdentityCredential()
token = cred.get_token("https://database.windows.net/.default").token

# pack token for pyodbc
token_bytes  = token.encode("utf-16-le")
token_struct = struct.pack(f"<I{len(token_bytes)}s", len(token_bytes), token_bytes)

conn = pyodbc.connect(
    "DRIVER={ODBC Driver 17 for SQL Server};"
    "SERVER=server.database.windows.net;"
    "DATABASE=mydb",
    attrs_before={1256: token_struct}   # SQL_COPT_SS_ACCESS_TOKEN = 1256
)

read - full table

df = (spark.read
      .jdbc(url=jdbc_url, table="dbo.orders", properties=connection_props))

read - query

df = (spark.read
      .jdbc(
          url        = jdbc_url,
          table      = "(SELECT * FROM dbo.orders WHERE dt >= '2024-01-01') AS t",
          properties = connection_props
      ))

read - parallel

manual

predicates = [
    "order_id BETWEEN 1       AND 250000",
    "order_id BETWEEN 250001  AND 500000",
    "order_id BETWEEN 500001  AND 750000",
    "order_id BETWEEN 750001  AND 1000000",
]

df = (spark.read
      .jdbc(url=jdbc_url, table="dbo.orders",
            predicates=predicates, properties=connection_props))

automatic

df = (spark.read
      .jdbc(
          url               = jdbc_url,
          table             = "(SELECT * FROM dbo.orders) AS t",
          column            = "order_id",        # numeric / date column
          lowerBound        = 1,
          upperBound        = 10_000_000,
          numPartitions     = 16,
          properties        = connection_props
      ))

write

append

(df.write
   .jdbc(url=jdbc_url, table="dbo.orders_out",
         mode="append", properties=connection_props))

overwrite

(df.write
   .option("truncate", "true")        # truncate instead of DROP + recreate
   .jdbc(url=jdbc_url, table="dbo.orders_out",
         mode="overwrite", properties=connection_props))

write - batch size + isolation

(df.write
   .option("batchsize",        10_000)
   .option("isolationLevel",   "READ_COMMITTED")
   .option("numPartitions",    8)              # parallel writers
   .jdbc(url=jdbc_url, table="dbo.orders_out",
         mode="append", properties=connection_props))

Gotchas

• JDBC reads are single-partition by default — always use numPartitions for large tables.
• lowerBound/upperBound are splitting hints not filters — rows outside range still read.
• overwrite without truncate=true drops and recreates table — kills indexes.
• batchsize default is 1000 — too low for large writes, set 10k–50k.
• Driver jar must be on cluster — in Databricks add com.microsoft.sqlserver:mssql-jdbc:12.4.2.jre11 as Maven library.
• Stored procs and DDL must go through a direct JDBC/pyodbc connection on the driver, not Spark.
• Never put credentials in code — always dbutils.secrets.get.
• Parallel writers (numPartitions on write) can cause deadlocks on SQL Server — tune based on target table's lock behaviour.

json

{"id": 1, "name": "alice", "address": {"city": "pune", "pin": "411001"}, "orders": [{"oid": 101, "amt": 500}, {"oid": 102, "amt": 300}]}

read

from pyspark.sql.types import *
from pyspark.sql.functions import col, explode, explode_outer

schema = StructType([
    StructField("id",      IntegerType()),
    StructField("name",    StringType()),
    StructField("address", StructType([
        StructField("city", StringType()),
        StructField("pin",  StringType())
    ])),
    StructField("orders",  ArrayType(StructType([
        StructField("oid", IntegerType()),
        StructField("amt", IntegerType())
    ])))
])

df = spark.read.schema(schema).json("/mnt/data/users.json")
df.show()

flatten and explode in 1 step

and also flatten exploded json

final = (df
    .select(
        "id",
        "name",
        col("address.city").alias("city"),
        col("address.pin").alias("pin"),
        explode("orders").alias("order")
    )
    .select(
        "id", "name", "city", "pin",
        col("order.oid").alias("order_id"),
        col("order.amt").alias("amount")
    ))

Gotchas

• explode on a null/empty array drops the row silently — use explode_outer to keep it.
• After explode, the array column is gone — select what you need before exploding or re-join.
• DDL schema shorthand for the same schema:

schema = "id INT, name STRING, address STRUCT<city:STRING, pin:STRING>, orders ARRAY<STRUCT<oid:INT, amt:INT>>"