from pyspark.sql import SparkSession
from pyspark.sql.functions import lit, col

# 创建SparkSession
spark = SparkSession.builder.appName("AddContentExample").getOrCreate()

# 示例数据
data = [("Alice", 25), ("Bob", 30), ("Charlie", 35)]
df = spark.createDataFrame(data, ["name", "age"])

# 方法1：使用withColumn添加新列
df_with_new_col = df.withColumn("country", lit("USA"))

# 方法2：基于现有列计算新列
df_with_calculated_col = df.withColumn("age_plus_10", col("age") + 10)

df_with_new_col.show()
df_with_calculated_col.show()

from pyspark.sql.functions import udf
from pyspark.sql.types import StringType

# 定义UDF
def add_title(name):
    return f"Mr/Ms {name}"

add_title_udf = udf(add_title, StringType())

# 应用UDF
df_with_title = df.withColumn("title_name", add_title_udf(col("name")))
df_with_title.show()

# 注册临时视图
df.createOrReplaceTempView("people")

# 使用SQL添加新列
df_sql = spark.sql("""
    SELECT *, 
           'USA' as country,
           age + 10 as age_plus_10,
           CONCAT('Mr/Ms ', name) as title_name
    FROM people
""")
df_sql.show()

# 转换为RDD并添加内容
rdd = df.rdd.map(lambda row: (row["name"], row["age"], "USA"))

# 转换回DataFrame
df_from_rdd = spark.createDataFrame(rdd, ["name", "age", "country"])
df_from_rdd.show()