PySpark DataFrame列拆分：提取trend前缀或首个下划线前内容

PySpark 实现代码 #

直接使用PySpark内置函数实现即可，无需自定义UDF，执行效率更高。

1. 导入依赖函数 #

from pyspark.sql import functions as F

2. 核心转换逻辑 #

# 替换下方df为你自己的DataFrame变量名即可
df_result = df.withColumn(
    "column_1",
    # 判断字段是否包含trend
    F.when(
        F.col("Column_1").contains("trend"),
        # 截取从开头到trend结束的内容，trend共5个字符，定位到首字符t的位置后+4即可取到完整trend
        F.substring(F.col("Column_1"), 1, F.locate("trend", F.col("Column_1")) + 4)
    ).otherwise(
        # 不包含trend则截取第一个下划线之前的内容
        F.substring(F.col("Column_1"), 1, F.locate("_", F.col("Column_1")) - 1)
    )
).select("column_1")

3. 完整测试示例（可直接运行验证） #

# 构造和你需求一致的测试数据
data = [
    ("daily_trend_navigator",),
    ("weekly_trend_navigator",),
    ("day_of_week_trend_display",),
    ("day_of_month_trend_notifier",),
    ("empty_navigator",),
    ("unique_notifier",)
]
df = spark.createDataFrame(data, schema=["Column_1"])

# 执行转换
df_result = df.withColumn(
    "column_1",
    F.when(
        F.col("Column_1").contains("trend"),
        F.substring(F.col("Column_1"), 1, F.locate("trend", F.col("Column_1")) + 4)
    ).otherwise(
        F.substring(F.col("Column_1"), 1, F.locate("_", F.col("Column_1")) - 1)
    )
).select("column_1")

# 输出结果
df_result.show(truncate=False)

运行输出 #

+-------------------+
|column_1           |
+-------------------+
|daily_trend        |
|weekly_trend       |
|day_of_week_trend  |
|day_of_month_trend |
|empty              |
|unique             |
+-------------------+

内容的提问来源于stack exchange，提问作者Padfoot123