使用PySpark查找包含特定值666的DataFrame列名

问题场景 #

给定如下PySpark DataFrame（包含字符串、数组、数值等混合类型）：

需要定位出包含数值666的列名，预期结果为col3。

解决方案 #

由于DataFrame存在多种数据类型，需要针对不同类型分别处理，判断列中是否存在匹配666的记录：

代码实现 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, array_contains, try_cast

# 初始化Spark会话
spark = SparkSession.builder.appName("FindTargetColumn").getOrCreate()

# 构建示例DataFrame
sample_data = [
    (1, [3,7], 5),
    ("hello", 4, 666),
    (4, "world", 4)
]
df = spark.createDataFrame(sample_data, schema=["col1", "col2", "col3"])

target_value = 666
matching_columns = []

# 遍历所有列逐一检查
for col_name in df.columns:
    col_type = df.schema[col_name].dataType.typeName()
    
    if col_type == "array":
        # 数组类型：检查是否包含目标值
        has_match = df.filter(array_contains(col(col_name), target_value)).count() > 0
    elif col_type == "string":
        # 字符串类型：尝试转整数后匹配，避免转换失败报错
        has_match = df.filter(try_cast(col(col_name), "int") == target_value).count() > 0
    elif col_type in ["int", "long", "float", "double"]:
        # 数值类型：直接相等判断
        has_match = df.filter(col(col_name) == target_value).count() > 0
    else:
        # 其他未覆盖类型暂时跳过
        continue
    
    if has_match:
        matching_columns.append(col_name)

# 输出结果
print(f"包含数值{target_value}的列：{matching_columns}")

代码说明 #

• 数组类型处理：用array_contains直接检查数组元素是否包含目标值。
• 字符串类型处理：借助try_cast安全转换为整数，避免非数字字符串导致的报错，再进行匹配。
• 数值类型处理：直接执行相等比较逻辑。
• 最后收集所有符合条件的列名，输出结果。

运行结果 #

执行代码后会输出：

包含数值666的列：['col3']

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