在pyspark中向数据帧添加唯一的连续行号

时间:2018-10-31 12:01:00

标签: csv dataframe pyspark rdd

我想在pyspark的数据框中添加唯一的行号,并且不想使用monotonicallyIncreasingId和partitionBy方法。 我认为这个问题可能是先前提出的类似问题的重复,无论我是否做对了,仍在寻找一些建议。 以下是我的代码段: 我有一个具有以下输入记录集的csv文件:

1,VIKRANT SINGH RANA    ,NOIDA   ,10000
3,GOVIND NIMBHAL        ,DWARKA  ,92000
2,RAGHVENDRA KUMAR GUPTA,GURGAON ,50000
4,ABHIJAN SINHA         ,SAKET   ,65000
5,SUPER DEVELOPER       ,USA     ,50000
6,RAJAT TYAGI           ,UP      ,65000
7,AJAY SHARMA           ,NOIDA   ,70000
8,SIDDHARTH BASU        ,SAKET   ,72000
9,ROBERT                ,GURGAON ,70000

并且我已经将此csv文件加载到数据框中。

PATH_TO_FILE="file:///u/user/vikrant/testdata/EMP_FILE.csv"

emp_df = spark.read.format("com.databricks.spark.csv") \
  .option("mode", "DROPMALFORMED") \
  .option("header", "true") \
  .option("inferschema", "true") \
  .option("delimiter", ",").load(PATH_TO_FILE)

+------+--------------------+--------+----------+
|emp_id|            emp_name|emp_city|emp_salary|
+------+--------------------+--------+----------+
|     1|VIKRANT SINGH RAN...|NOIDA   |     10000|
|     3|GOVIND NIMBHAL   ...|DWARKA  |     92000|
|     2|RAGHVENDRA KUMAR ...|GURGAON |     50000|
|     4|ABHIJAN SINHA    ...|SAKET   |     65000|
|     5|SUPER DEVELOPER  ...|USA     |     50000|
|     6|RAJAT TYAGI      ...|UP      |     65000|
|     7|AJAY SHARMA      ...|NOIDA   |     70000|
|     8|SIDDHARTH BASU   ...|SAKET   |     72000|
|     9|ROBERT           ...|GURGAON |     70000|
+------+--------------------+--------+----------+

empRDD = emp_df.rdd.zipWithIndex()
newRDD=empRDD.map(lambda x: (list(x[0]) + [x[1]]))
 newRDD.take(2);
[[1, u'VIKRANT SINGH RANA    ', u'NOIDA   ', 10000, 0], [3, u'GOVIND NIMBHAL        ', u'DWARKA  ', 92000, 1]]

当我将int值添加到列表中时,我丢失了数据框架构。

newdf=newRDD.toDF(['emp_id','emp_name','emp_city','emp_salary','row_id'])
newdf.show();

+------+--------------------+--------+----------+------+
|emp_id|            emp_name|emp_city|emp_salary|row_id|
+------+--------------------+--------+----------+------+
|     1|VIKRANT SINGH RAN...|NOIDA   |     10000|     0|
|     3|GOVIND NIMBHAL   ...|DWARKA  |     92000|     1|
|     2|RAGHVENDRA KUMAR ...|GURGAON |     50000|     2|
|     4|ABHIJAN SINHA    ...|SAKET   |     65000|     3|
|     5|SUPER DEVELOPER  ...|USA     |     50000|     4|
|     6|RAJAT TYAGI      ...|UP      |     65000|     5|
|     7|AJAY SHARMA      ...|NOIDA   |     70000|     6|
|     8|SIDDHARTH BASU   ...|SAKET   |     72000|     7|
|     9|ROBERT           ...|GURGAON |     70000|     8|
+------+--------------------+--------+----------+------+

我做对了吗?还是有更好的方法在pyspark中添加或保留数据框的架构?

使用zipWithIndex方法为大型数据帧添加唯一的连续行号是否可行?我们可以使用该row_id重新划分数据帧以在分区之间均匀分布数据吗?

3 个答案:

答案 0 :(得分:3)

我找到了一个解决方案,这很简单。 由于我的数据框中没有在所有行中都具有相同值的列,因此在与partitionBy子句一起使用时,使用row_number不会生成唯一的行号。

让我们在现有数据框中添加一个带有默认值的新列。

emp_df= emp_df.withColumn("new_column",lit("ABC"))

并使用“ new_column”列创建具有parition的窗口函数

w = Window().partitionBy('new_column').orderBy(lit('A'))
df = emp_df.withColumn("row_num", row_number().over(w)).drop("new_column")

您将获得理想的结果:

+------+--------------------+--------+----------+-------+
|emp_id|            emp_name|emp_city|emp_salary|row_num|
+------+--------------------+--------+----------+-------+
|     1|VIKRANT SINGH RAN...|NOIDA   |     10000|      1|
|     2|RAGHVENDRA KUMAR ...|GURGAON |     50000|      2|
|     7|AJAY SHARMA      ...|NOIDA   |     70000|      3|
|     9|ROBERT           ...|GURGAON |     70000|      4|
|     4|ABHIJAN SINHA    ...|SAKET   |     65000|      5|
|     8|SIDDHARTH BASU   ...|SAKET   |     72000|      6|
|     5|SUPER DEVELOPER  ...|USA     |     50000|      7|
|     3|GOVIND NIMBHAL   ...|DWARKA  |     92000|      8|
|     6|RAJAT TYAGI      ...|UP      |     65000|      9|
+------+--------------------+--------+----------+-------+

答案 1 :(得分:1)

由于我的声誉,我无法编辑我的问题或答案,因此以下是我对使用窗口函数为Pyspark中给定数据帧生成唯一行号的想法。

我尝试在emp_city列上将row_number与窗口函数和partitionBy一起使用。由于emp_city列的值不相同,因此不会生成唯一的行号。 如何在数据框中获得n行的唯一连续行号。我能够使用zipWithIndex做到这一点。

from pyspark.sql.functions import row_number,lit,rank
from pyspark.sql.window import Window

w = Window().partitionBy('emp_city').orderBy(lit('A'))
df = emp_df.withColumn("row_num", row_number().over(w))

+------+--------------------+--------+----------+-------+
|emp_id|            emp_name|emp_city|emp_salary|row_num|
+------+--------------------+--------+----------+-------+
|     3|GOVIND NIMBHAL   ...|DWARKA  |     92000|      1|
|     4|ABHIJAN SINHA    ...|SAKET   |     65000|      1|
|     8|SIDDHARTH BASU   ...|SAKET   |     72000|      2|
|     1|VIKRANT SINGH RAN...|NOIDA   |     10000|      1|
|     7|AJAY SHARMA      ...|NOIDA   |     70000|      2|
|     2|RAGHVENDRA KUMAR ...|GURGAON |     50000|      1|
|     9|ROBERT           ...|GURGAON |     70000|      2|
|     6|RAJAT TYAGI      ...|UP      |     65000|      1|
|     5|SUPER DEVELOPER  ...|USA     |     50000|      1|
+------+--------------------+--------+----------+-------+

如果我不使用partitionBy,则按需使用window函数可以正常工作,但是它将所有数据从“ n”个分区移至一个分区

newdf=emp_df.repartition("emp_city")
partitionSizes = newdf.rdd.glom().map(len).collect();
print partitionSizes
[0, 0, 0, 0, 4, 0, 2, 1, 1, 1]
w = Window().orderBy(lit('A'))
df = newdf.withColumn("row_num", row_number().over(w))

18/11/01 05:46:12 WARN WindowExec: No Partition Defined for Window operation! Moving all data to a single partition, this can cause serious performance degradation.
+------+--------------------+--------+----------+-------+
|emp_id|            emp_name|emp_city|emp_salary|row_num|
+------+--------------------+--------+----------+-------+
|     1|VIKRANT SINGH RAN...|NOIDA   |     10000|      1|
|     2|RAGHVENDRA KUMAR ...|GURGAON |     50000|      2|
|     7|AJAY SHARMA      ...|NOIDA   |     70000|      3|
|     9|ROBERT           ...|GURGAON |     70000|      4|
|     4|ABHIJAN SINHA    ...|SAKET   |     65000|      5|
|     8|SIDDHARTH BASU   ...|SAKET   |     72000|      6|
|     5|SUPER DEVELOPER  ...|USA     |     50000|      7|
|     3|GOVIND NIMBHAL   ...|DWARKA  |     92000|      8|
|     6|RAJAT TYAGI      ...|UP      |     65000|      9|
+------+--------------------+--------+----------+-------+

我正在寻找一些可行的解决方案,以为我的数据框生成一个新的唯一编号,该编号应该是连续的。

答案 2 :(得分:0)

使用 Spark SQL:

df = spark.sql("""
SELECT 
    row_number() OVER (
        PARTITION BY '' 
        ORDER BY '' 
    ) as id,
    *
FROM 
    VALUES 
    ('Bob  ', 20),
    ('Alice', 21),
    ('Gary ', 21),
    ('Kent ', 25),
    ('Gary ', 35)
""")

输出:

>>> df.printSchema()
root
 |-- id: integer (nullable = true)
 |-- col1: string (nullable = false)
 |-- col2: integer (nullable = false)

>>> df.show()
+---+-----+----+
| id| col1|col2|
+---+-----+----+
|  1|Bob  |  20|
|  2|Alice|  21|
|  3|Gary |  21|
|  4|Kent |  25|
|  5|Gary |  35|
+---+-----+----+