DataFrame API 复杂查询和聚合操作

1. 基础查询操作

选择与过滤：

# 列选择
df.select("name", "age", "salary")
df.select(col("name"), col("age").alias("用户年龄"))

# 条件过滤
df.filter(col("age") > 18)
df.filter((col("salary") > 5000) & (col("dept") == "IT"))

排序与去重：

df.orderBy(col("salary").desc(), col("age").asc())
df.distinct()
df.dropDuplicates(["name", "dept"])

2. 复杂聚合操作

基本聚合：

from pyspark.sql.functions import *

# 单维度聚合
df.groupBy("dept").agg(
    avg("salary").alias("avg_salary"),
    sum("salary").alias("total_salary"),
    count("*").alias("employee_count")
)

多维度聚合（数据立方体）：

# 多维分组聚合
df.cube("dept", "gender").agg(
    avg("salary").alias("avg_salary"),
    count("*").alias("count")
)

# 滚动窗口聚合
df.rollup("year", "month").agg(sum("sales"))

窗口函数（高级分析）：

from pyspark.sql.window import Window

# 定义窗口规范
window_spec = Window.partitionBy("dept").orderBy("salary")

df.withColumn("rank", rank().over(window_spec)) \
  .withColumn("salary_diff", col("salary") - lag("salary", 1).over(window_spec)) \
  .withColumn("cumulative_salary", sum("salary").over(window_spec))

3. 多表连接操作

各种连接类型：

# 内连接
df1.join(df2, df1.id == df2.id, "inner")

# 左外连接
df1.join(df2, df1.id == df2.id, "left_outer")

# 全外连接
df1.join(df2, df1.id == df2.id, "full_outer")

# 交叉连接
df1.crossJoin(df2)

4. 高级查询技巧

条件表达式：

from pyspark.sql.functions import when

df.withColumn("salary_level", 
    when(col("salary") > 10000, "高薪")
     .when(col("salary") > 5000, "中薪")
     .otherwise("低薪")
)

JSON数据解析：

# 解析JSON字符串列
df.withColumn("json_data", from_json(col("json_string"), schema)) \
  .select("json_data.*")

数组和Map操作：

# 数组操作
df.withColumn("array_size", size(col("tags"))) \
  .withColumn("first_tag", element_at(col("tags"), 1))

# Map操作
df.withColumn("map_keys", map_keys(col("properties"))) \
  .withColumn("map_values", map_values(col("properties")))

5. 性能优化技巧

分区和缓存：

# 重新分区
df.repartition(100, col("dept"))

# 缓存常用DataFrame
df.cache()

# 广播小表
from pyspark.sql.functions import broadcast
df1.join(broadcast(df2), "id")

查询执行计划查看：

df.explain()  # 查看物理执行计划
df.explain("extended")  # 查看详细执行计划

6. 实际应用示例

销售数据分析：

sales_df.groupBy("product_category", "region") \
    .agg(
        sum("revenue").alias("total_revenue"),
        avg("unit_price").alias("avg_price"),
        countDistinct("customer_id").alias("unique_customers")
    ) \
    .filter(col("total_revenue") > 100000) \
    .orderBy(col("total_revenue").desc())

用户行为分析：

window_spec = Window.partitionBy("user_id").orderBy("timestamp")

user_behavior_df \
    .withColumn("prev_action", lag("action", 1).over(window_spec)) \
    .withColumn("session_gap", 
                unix_timestamp("timestamp") - unix_timestamp(lag("timestamp", 1).over(window_spec))) \
    .withColumn("new_session", 
                when((col("prev_action").isNull()) | (col("session_gap") > 1800), 1).otherwise(0)) \
    .withColumn("session_id", sum("new_session").over(window_spec))

这些操作组合使用可以处理绝大多数复杂的数据分析场景，同时保持Spark的分布式计算优势。