在指定位置插入列

Question

如何在指定位置插入列而不列出所有现有列名？

我有这个数据框：

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'),
)
df.show()
# +---+---+---+
# |  a|  b|  c|
# +---+---+---+
# | 11| 22| 33|
# +---+---+---+

我想在列之后插入一列

a

。

我知道我可以按照我喜欢的顺序列出所有列：

df.select('a', F.lit('foo'), 'b', 'c')

但是我想插入一列而不列出现有的列名称。

Answer 1

这样的东西可以工作：

position = 1
cols = df.columns
df = df.select(*cols[:position], F.lit('foo'), *cols[position:])

df.show()
# +---+---+---+---+
# |  a|foo|  b|  c|
# +---+---+---+---+
# | 11|foo| 22| 33|
# +---+---+---+---+

---

可以创建插入或移动列的函数：

def insert_col(df, position, col_name, col_expr=None):
    cols = [c for c in df.columns if c != col_name]
    col = F.col(col_name) if col_expr is None else col_expr
    df = df.select(*cols[:position], col.alias(col_name), *cols[position:])
    return df

df = insert_col(df, 1, 'foo', F.lit('bar'))
df.show()
# +---+---+---+---+
# |  a|foo|  b|  c|
# +---+---+---+---+
# | 11|bar| 22| 33|
# +---+---+---+---+

df = insert_col(df, 0, 'c')
df.show()
# +---+---+---+
# |  c|  a|  b|
# +---+---+---+
# | 33| 11| 22|
# +---+---+---+