文章目录

Spark DataFrame 添加索引列的三种方法方法一方法二方法三

Spark DataFrame 添加索引列的三种方法

刚开始用Spark，操作dataframe不是很熟练，遇到的第一个问题是给dataframe添加索引列，查阅了网上的一些教程，大都是用Scala语言编写的代码，下面给出自己用python写的三种方法。

方法一：先创建Pandas版本的dataframe，然后带索引保存为本地文件，再用SparkSession来创建DataFrame。方法二：先创建Pandas版本的dataframe，添加一个索引列后，再直接转换成Spark版本的dataframe。方法三：直接利用Spark的select方法，新增列的过程包含在自定义的udf函数中。

---

方法一

1.先创建一个Pandas版本的dataframe，从本地csv文件导入数据。

import pandas as pd from pyspark.sql import * # build a dataframe pandas_df = pd.read_csv('/home/hadoop/PycharmProjects/wjw/ftdd/ftdd/data/net_data.csv') print(pandas_df)

显示结果如下 2.然后添加索引列，从1开始编号，并保存为本地文件。

pandas_df.index = range(1, len(pandas_df) + 1) # 索引值设置为从1开始 print(pandas_df) pandas_df.to_csv('/home/hadoop/PycharmProjects/wjw/ftdd/ftdd/data/net_data01.csv', index=True, index_label='index')

显示结果 3.再用SparkSession来创建DataFrame。

spark = SparkSession.builder.getOrCreate() df = spark.read.csv('/home/hadoop/PycharmProjects/wjw/ftdd/ftdd/data/net_data01.csv', header=True) df.show()

结果如下

---

方法二

1.第一步同方法一。 2.先添加一个索引列，insert方法可以设置列的插入位置。

pandas_df.index = range(1, len(pandas_df) + 1) # set index from 1 pandas_index = pandas_df.index # 将dataframe的索引赋给一个变量 pandas_df.insert(0, 'index', pandas_index) # 第一个参数是列插入的位置 print(pandas_df)

结果如下 3. 再直接转换成Spark版本的dataframe

spark = SparkSession.builder.getOrCreate() spark_df = spark.createDataFrame(pandas_df) spark_df.show()

结果如下

---

方法三

1.用SparkSession来创建dataframe，从本地直接读取数据。

from pyspark.sql import SparkSession from pyspark.sql.functions import udf, col from pyspark.sql.types import IntegerType spark = SparkSession.builder.getOrCreate() spark_df = spark.read.csv('/home/hadoop/PycharmProjects/wjw/ftdd/ftdd/data/net_data.csv', header=True) spark_df.show()

结果如下 2. 定义自己的udf函数，来创建索引列。

index_list = [x for x in range(1, spark_df.count()+1)] # 构造一个列表存储索引值，用生成器会出错 idx = 0 # 定义一个函数 def set_index(x): global idx # 将idx设置为全局变量 if x is not None: idx += 1 return index_list[idx-1]

3.调用select方法来添加索引列。

index = udf(set_index, IntegerType()) # udf的注册，这里需要定义其返回值类型 spark_df.select(col("*"), index("cab_id").alias("index")).show() # udf的注册的使用，alias方法用于修改列名

结果如下