python - 如何根据行的内容分割pyspark数据帧

Question

我想根据 Pyspark DataFrame 中行的第一个字符拆分文件。

原始数据有一列，数据包括

1. 文件名(例如“DATE20191009”)

2. 文件内容(例如“1”、“2”、“3”)

输入示例文件(Pyspark DataFrame):

column1

Date20191009

1

2

3

Date20191010

1

4

5

我想获取一个 Pyspark DataFrame，其文件名作为数据的分割。

文件名放置在DataFrame的column1中，文件内容放置在DataFrame的column2中。

预期输出(Pyspark DataFrame)

column1  column2
Date20191009 [1,2,3]
Date20191010 [1,4,5]

我尝试了Pandas Dataframe和Pyspark DataFrame.collect()，但都因数据量过多(超过900万行)而失败。

Answer 1

>>> from pyspark.sql.window import Window
>>> from pyspark.sql.functions import *
>>> w = Window.rowsBetween(Window.unboundedPreceding, 0)

   #Input DataFrame

>>> df.show()
+------------+
|     column1|
+------------+
|Date20191009|
|           1|
|           2|
|           3|
|Date20191010|
|           1|
|           4|
|           5|
+------------+

>>> df1 = df.withColumn('tmp', when(df.column1.startswith('Date'), df.column1).otherwise(None)).withColumn('temp', last('tmp', True).over(w)).drop('tmp')
>>> df1.show()

+------------+------------+
|     column1|        temp|
+------------+------------+
|Date20191009|Date20191009|
|           1|Date20191009|
|           2|Date20191009|
|           3|Date20191009|
|Date20191010|Date20191010|
|           1|Date20191010|
|           4|Date20191010|
|           5|Date20191010|
+------------+------------+

>>> df1.filter(df1.column1 != df1.temp).groupBy(df1.temp).agg(concat_ws(',',collect_list(df1.column1)).alias('column2')).withColumnRenamed("temp", "column1").show()

+------------+-------+
|     column1|column2|
+------------+-------+
|Date20191009|  1,2,3|
|Date20191010|  1,4,5|
+------------+-------+