Spark withColumn基于日期的顺序整数

时间:2017-09-08 21:03:56

标签: python apache-spark pyspark spark-dataframe

我有一个数据集,其公司的每日产品生产率。我想在该数据框中添加一个列,该列将按日期顺序编号。

实施例
Acme Product1 1/1/2000 5
Acme Product1 1/2/2000 7
Acme Product2 3/1/2000 9
Acme Product2 3/2/2000 4
公司2产品X 4/1/2015 6
公司2产品X 4/2/2015 3

我想添加一个新列,如:
Acme Product1 1/1/2000 5 1
Acme Product1 1/2/2000 7 2
Acme Product2 3/1/2000 9 1
Acme Product2 3/2/2000 4 2
公司2产品X 4/1/2015 6 1
公司2产品X 4/2/2015 3 2
公司2产品X 4/2/2015 2 3

这一切都是为了让我可以根据新专栏对公司及其产品进行比较。因此,无论日期如何,他们所有的一天生产都是产品。

1 个答案:

答案 0 :(得分:0)

您可以使用pyspark.sql函数row_number进行窗口化(别名rowNumber适用于spark <= 1.6.X)。

首先让我们创建数据框:

myDF = spark.createDataFrame(
    sc.parallelize([["Acme", "Product1", "1/1/2000", 5],
        ["Acme", "Product1", "1/2/2000", 7],
        ["Acme", "Product2", "3/1/2000", 9],
        ["Acme", "Product2", "3/2/2000", 4],
        ["Company2", "ProductX", "4/1/2015", 6],
        ["Company2", "ProductX", "4/2/2015", 3], 
        ["Company2", "ProductX", "4/2/2015", 2]]), 
    ["company", "product", "date", "nb"])

    +--------+--------+--------+---+
    | company| product|    date| nb|
    +--------+--------+--------+---+
    |    Acme|Product1|1/1/2000|  5|
    |    Acme|Product1|1/2/2000|  7|
    |    Acme|Product2|3/1/2000|  9|
    |    Acme|Product2|3/2/2000|  4|
    |Company2|ProductX|4/1/2015|  6|
    |Company2|ProductX|4/2/2015|  3|
    |Company2|ProductX|4/2/2015|  2|
    +--------+--------+--------+---+

现在使用窗口函数:

from pyspark.sql.window import Window
import pyspark.sql.functions as F

resultDF = myDF.withColumn(
    "rowNum", 
    F.row_number().over(Window.partitionBy("company", "product")
                           .orderBy("date")))

    +--------+--------+--------+---+------+
    | company| product|    date| nb|rowNum|
    +--------+--------+--------+---+------+
    |Company2|ProductX|4/1/2015|  6|     1|
    |Company2|ProductX|4/2/2015|  3|     2|
    |Company2|ProductX|4/2/2015|  2|     3|
    |    Acme|Product2|3/1/2000|  9|     1|
    |    Acme|Product2|3/2/2000|  4|     2|
    |    Acme|Product1|1/1/2000|  5|     1|
    |    Acme|Product1|1/2/2000|  7|     2|
    +--------+--------+--------+---+------+