PySpark替代Pandas iloc函数报错TypeError的解决方法求助

错误原因 #

触发的TypeError: list indices must be integer or slices, not list，是因为x.columns是PySpark DataFrame的列名列表，而你用列表类型的v执行切片操作（v:），列表的索引/切片仅支持整数或slice对象，不支持列表，因此报错。

原Pandas代码逻辑解析 #

原代码x.iloc[1:, v]的作用是：

• 跳过x的第一行（iloc[1:]）
• 选择v指定的列（v可能是单个列索引、列索引列表、列名或列名列表）

分场景解决方案 #

场景1：v是单个列索引（整数） #

from pyspark.sql.functions import monotonically_increasing_id

for k, v in df.items():
    # 给DataFrame添加行号列，用于跳过第一行
    x_with_row_num = x.withColumn("row_num", monotonically_increasing_id())
    # 筛选行号>0（对应iloc[1:]），并选择指定索引的列
    xy = x_with_row_num.filter(x_with_row_num.row_num > 0).select(x.columns[v])

场景2：v是列索引列表 #

from pyspark.sql.functions import monotonically_increasing_id

for k, v in df.items():
    # 根据索引列表提取对应的列名
    selected_cols = [x.columns[idx] for idx in v]
    x_with_row_num = x.withColumn("row_num", monotonically_increasing_id())
    # 筛选行并选择目标列
    xy = x_with_row_num.filter(x_with_row_num.row_num > 0).select(*selected_cols)

场景3：v是列名（字符串）或列名列表 #

from pyspark.sql.functions import monotonically_increasing_id

for k, v in df.items():
    x_with_row_num = x.withColumn("row_num", monotonically_increasing_id())
    # 处理单个列名或列名列表的情况
    if isinstance(v, str):
        xy = x_with_row_num.filter(x_with_row_num.row_num > 0).select(v)
    else:
        xy = x_with_row_num.filter(x_with_row_num.row_num > 0).select(*v)

注意事项 #

PySpark没有像Pandas那样的行位置索引（iloc），monotonically_increasing_id()生成的行号是全局唯一递增的，但如果数据有分区，行号不一定连续。如果需要严格按原始顺序跳过第一行，需先对数据基于主键排序后再添加行号。

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