PySpark:迭代DataFrame中的小组

时间:2016-11-10 06:29:00

标签: python pyspark grouping spark-dataframe

我试图了解如何在PySpark DataFrame中的小组内进行操作。假设我有以下架构的DF:

root
|-- first_id: string (nullable = true)
|-- second_id_struct: struct (nullable = true)
|    |-- s_id: string (nullable = true)
|    |-- s_id_2: int (nullable = true)
|-- depth_from: float (nullable = true)
|-- depth_to: float (nullable = true)
|-- total_depth: float (nullable = true) 

因此数据可能如下所示:

我想:

  1. 按first_id分组数据
  2. 在每个组内,按s_id_2按升序排序
  3. 将额外列layer附加到结构或根DataFrame,以指示组中此s_id_2的顺序。
  4. 例如:

    first_id | second_id | second_id_order 
    ---------| --------- | ---------------
          A1 |   [B, 10] | 1  
    ---------| --------- | ---------------
          A1 |   [B, 14] | 2
    ---------| --------- | ---------------
          A1 |   [B, 22] | 3
    ---------| --------- | ---------------
          A5 |    [A, 1] | 1
    ---------| --------- | ---------------
          A5 |    [A, 7] | 2
    ---------| --------- | ---------------
          A7 |      null | 1
    ---------| --------- | ---------------        
    

    一旦分组,每个first_id最多只有4 second_id_struct。我该如何处理这类问题?

    我对如何在一般的小组(1-40行)DataFrame中进行迭代操作特别感兴趣,其中组内的列顺序很重要。

    谢谢!

1 个答案:

答案 0 :(得分:1)

创建一个DataFrame

to_timestamp('12-JUN-99','DD-MM-YY HH.MI.SSXFF AM')

你可以看到结构

d = [{'first_id': 'A1', 'second_id': ['B',10]}, {'first_id': 'A1', 'second_id': ['B',14]},{'first_id': 'A1', 'second_id': ['B',22]},{'first_id': 'A5', 'second_id': ['A',1]},{'first_id': 'A5', 'second_id': ['A',7]}]

df = sqlContext.createDataFrame(d) 

然后您可以使用 dense_rank Window 功能来显示子组中的顺序。它与SQL中的over partition相同。

窗口功能的介绍:Introducing Window Functions in Spark SQL

代码在这里:

df.printSchema()

|-- first_id: string (nullable = true)
|-- second_id: array (nullable = true)
|........|-- element: string (containsNull = true)

df.show()
+--------+----------+
|first_id|second_id |
+--------+----------+
|      A1|   [B, 10]|
|      A1|   [B, 14]|
|      A1|   [B, 22]|
|      A5|    [A, 1]|
|      A5|    [A, 7]|
+--------+----------+

结果:

# setting a window spec
windowSpec = Window.partitionBy('first_id').orderBy(df.second_id[1])
# apply dense_rank to the window spec
df.select(df.first_id, df.second_id, dense_rank().over(windowSpec).alias("second_id_order")).show()