pyspark dataframe将一行分成多行并标记序号(index)

最新推荐文章于 2024-05-15 00:49:49 发布

原创

最新推荐文章于 2024-05-15 00:49:49 发布 · 5.9k 阅读

2 ·

CC 4.0 BY-SA版权

文章标签：

#spark

本文介绍了如何利用pyspark.sql.functions的udf和explode函数，将dataframe的一行数据拆分为多行，并在转换过程中添加字典序作为序号。在操作中需要注意udf返回的数据类型必须是map，否则explode操作会因为数据类型不匹配而报错。

原始数据如下：

gid	score
a1	90 80 79 80
a2	79 89 45 60
a3	57 56 89 75

from pyspark.sql.functions import udf, col
from pyspark.sql.types import MapType, IntegerType, StringType

def udf_array_to_map(array):
    if array is None:
        return array
    return dict((i, v) for i, v in enumerate(array))

# col(): returns a column based on the given column name
# MapType: 表示包括一组key-value的值．通过keyType表示key数据的类型，通过valueType表示value数据的类型．
#          最后一个参数指明mapType重点值是否有null值
def generate_idx_for_df(df, id_name, col_name, col_schema):
    """
    generate_idx_for_df, explodes rows with array as a column into a new row for each
    element in the array, with 'INTEGER_IDX' indicating its index in the original array.
    :param df: dataframe w