使用PySp处理配置单元中每个组的数据

Question

我有一些用户数据存储在配置单元表（ORC文件格式）中，每个用户有多行。例如：

user | attr1 | attr2
---- | ----- | -----
u1   | 1     | a
u1   | 2     | b
u2   | 3     | a
u2   | 4     | b

我需要映射属于同一用户的每一组行来计算一些字符串。在本例中，f([(1, a), (2, b)])和f([(3, a), (4, b)])。在

如何使用Hive+Spark的pythonapi实现这一点？在

据我所见，PythonGroupedDataAPI只有stat函数。我应该减少底层行RDD吗？有没有一种更有效的方式可以利用Hive/ORC列格式进行分组？在

Answer 1

from pyspark.sql import functions as f

df = spark.createDataFrame([['u1', 1, 'a'], ['u1', 2, 'b'], ['u2', 3, 'a'], ['u2', 4, 'b']],
                           schema=['user', 'attr1', 'attr2'])
df = df.groupBy('user').agg(f.collect_list(f.array(df['attr1'], df['attr2'])).alias('array_col'))
df = df.withColumn('len_col', f.size(df['new_col']))
df.show(truncate=False)

+  +                    +   -+                         
|user|array_col                               |len_col|
+  +                    +   -+
|u1  |[WrappedArray(1, a), WrappedArray(2, b)]|2      |
|u2  |[WrappedArray(3, a), WrappedArray(4, b)]|2      |
+  +                    +   -+