from pyspark.sql import SparkSession

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

df = spark.read.csv("sales.csv", header=True, inferSchema=True)

# DataFrame API
df.groupBy("region").sum("amount").show()

# Or pure SQL — same engine, same optimizer
df.createOrReplaceTempView("sales")
spark.sql("SELECT region, SUM(amount) FROM sales GROUP BY region").show()

# Read a continuous stream from Kafka
stream_df = spark.readStream \
    .format("kafka") \
    .option("kafka.bootstrap.servers", "host:9092") \
    .option("subscribe", "orders") \
    .load()

# Same DataFrame API as batch!
query = stream_df.groupBy("region").count() \
    .writeStream \
    .format("console") \
    .outputMode("complete") \
    .start()

from pyspark.ml.feature import VectorAssembler
from pyspark.ml.regression import LinearRegression

# Combine feature columns into a single vector column
assembler = VectorAssembler(inputCols=["age", "salary"], outputCol="features")
data = assembler.transform(df)

# Train a distributed linear regression model
lr = LinearRegression(featuresCol="features", labelCol="spend")
model = lr.fit(data)

# Cluster mode — driver runs INSIDE the cluster (on a NodeManager)
$ spark-submit --master yarn --deploy-mode cluster my_job.py

# Client mode — driver runs on the machine you submit from
$ spark-submit --master yarn --deploy-mode client my_job.py

$ spark-submit \
    --master k8s://https://<k8s-api-server>:443 \
    --deploy-mode cluster \
    --conf spark.kubernetes.container.image=my-spark-image:latest \
    --conf spark.executor.instances=5 \
    local:///opt/spark/jobs/my_job.py

from pyspark.sql import SparkSession

spark = SparkSession.builder \
    .appName("OrderProcessing") \
    .master("local[4]") \
    .config("spark.executor.memory", "4g") \
    .getOrCreate()

# Now you can use spark for everything:
df = spark.read.csv("data.csv")         # DataFrame API
spark.sql("SELECT * FROM my_view")       # Spark SQL
spark.readStream.format("kafka")         # Streaming
spark.conf.get("spark.sql.shuffle.partitions")  # Config

df = spark.read.csv("orders.csv", header=True)

# These are TRANSFORMATIONS — lazy, just added to the DAG, nothing runs yet
filtered = df.filter(df.amount > 100)
grouped  = filtered.groupBy("region").sum("amount")
sorted_  = grouped.orderBy("sum(amount)", ascending=False)

# This is an ACTION — NOW Spark builds the DAG and executes it
sorted_.show()

df = spark.read.csv("orders.csv", header=True, inferSchema=True)  # 8 partitions

# filter() and select() = narrow transformations, no shuffle → STAGE 1 (8 tasks)
filtered = df.filter(df.amount > 100).select("region", "amount")

# groupBy() = wide transformation, causes a SHUFFLE → new STAGE
result = filtered.groupBy("region").sum("amount")

# .show() = ACTION → creates JOB 1, which has 2 stages
#   Job 1
#     ├── Stage 1: read + filter + select  → 8 tasks (one per input partition)
#     └── Stage 2: groupBy + sum (after shuffle) → 200 tasks (default shuffle partitions)
result.show()

# .write() = ACTION → creates JOB 2 (separate job!)
result.write.parquet("output/")

$ spark-submit \
    --master yarn \
    --deploy-mode cluster \
    --driver-memory 4g \
    --executor-memory 8g \
    --executor-cores 4 \
    --num-executors 10 \
    my_etl_job.py

# Total cluster resources used:
# Executors: 10 x (8GB + 4 cores) = 80GB RAM, 40 cores
# Plus driver: 4GB RAM, 1 core

# Running locally in VS Code, but executing on a Databricks cluster
from databricks.connect import DatabricksSession

spark = DatabricksSession.builder.profile("my-profile").getOrCreate()

# df.show() runs on the remote Databricks cluster,
# but you can set breakpoints and debug locally!
df = spark.read.table("main.sales.orders")
df.show()