如何在PySpark DataFrame指定位置插入列且无需列出所有现有列名？

可以通过操作列名列表来实现需求，具体步骤和代码如下：

实现思路 #

1. 获取当前DataFrame的完整列名列表
2. 定位到目标列的索引位置
3. 构造包含新列的列顺序列表
4. 基于新列顺序重新选取数据

示例代码 #

from pyspark.sql import functions as F

df = spark.range(1).select(
    F.lit(11).alias('a'),
    F.lit(22).alias('b'),
    F.lit(33).alias('c'),
)

# 获取所有列名
cols = df.columns
# 找到目标列'a'的索引位置
target_idx = cols.index('a')
# 构造新的列序列：原列到a为止 + 新列 + a之后的剩余列
new_cols = cols[:target_idx+1] + [F.lit('foo').alias('new_col')] + cols[target_idx+1:]
# 生成新的DataFrame
df_new = df.select(new_cols)
df_new.show()

运行后输出结果：

+---+-------+---+---+
|  a|new_col|  b|  c|
+---+-------+---+---+
| 11|    foo| 22| 33|
+---+-------+---+---+

补充提示 #

• 如果需要避免目标列不存在时报错，可以先做列存在性判断：

  if 'a' in cols:
      target_idx = cols.index('a')
      # 后续构造列序列的操作
  else:
      print("目标列不存在，请检查列名")
  

• 新列不限于常量列，还可以是计算列、UDF生成列等任意Spark支持的列表达式。

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