PySpark：从数据帧创建字典的字典？

我有以下格式的数据，这些数据是从 Hive 获取到数据帧中的：

date, stock, price
1388534400, GOOG, 50
1388534400, FB, 60
1388534400, MSFT, 55
1388620800, GOOG, 52
1388620800, FB, 61
1388620800, MSFT, 55

其中 date 是当天午夜的纪元，我们有 10 年前左右的数据（8 亿多行）。 我想要一本字典如下：

{
'GOOG':
{
'1388534400': 50,
'1388620800': 52
}

'FB':
{
'1388534400': 60,
'1388620800': 61
}
}

一种天真的方法是获取唯一股票的列表，然后通过仅筛选出每只股票的那些行来获取数据帧的子集，但这似乎过于天真且效率极低。 这在 Spark 中可以轻松完成吗？我目前已经使用 PyHive 在本机 Python 中工作，但由于数据量巨大，我宁愿在集群/Spark 上完成此操作。

在spark 2.4中你可以使用map_from_arrays在汇总每只股票的价值时构建日期价值地图。那么就只是使用的问题了create_map使用股票代码作为键。本示例使用ChainMap从 python 3.4 开始构建最终的 dict 结构，如您所描述的。

import json
from collections import ChainMap
from pyspark.sql import SparkSession
from pyspark.sql.functions import *

spark = SparkSession \
    .builder \
    .appName("example") \
    .getOrCreate()

df = spark.createDataFrame([
    (1388534400, "GOOG", 50),
    (1388534400, "FB", 60),
    (1388534400, "MSFT", 55),
    (1388620800, "GOOG", 52),
    (1388620800, "FB", 61),
    (1388620800, "MSFT", 55)]
).toDF("date", "stock", "price")

out = df.groupBy("stock") \
        .agg(
            map_from_arrays(
                collect_list("date"), collect_list("price")).alias("values")) \
        .select(create_map("stock", "values").alias("values")) \
        .rdd.flatMap(lambda x: x) \
        .collect()

print(json.dumps(dict(ChainMap(*out)), indent=4, separators=(',', ': '), sort_keys=True))

这使：

{                                                                               
    "FB": {
        "1388534400": 60,
        "1388620800": 61
    },
    "GOOG": {
        "1388534400": 50,
        "1388620800": 52
    },
    "MSFT": {
        "1388534400": 55,
        "1388620800": 55
    }
}

However，正如你所说，你有很多数据，你可能实际上不想在内存中创建这个字典，所以你最好将其拆分并将相同的字典结构写入不同分区的文件中。

让我们通过将日期截断为给定月份并为每个月和每只股票编写单独的文件来实现这一点：

out = df.groupBy(trunc(expr("CAST(date as TIMESTAMP)"), "month").alias("month"), df["stock"]) \
        .agg(
            map_from_arrays(
                collect_list("date"), collect_list("price")).alias("values")) \
        .select("month", "stock", create_map("stock", "values").alias("values"))

out.write.partitionBy("month", "stock").format("json").save("out/prices")

这将为您提供如下结构：

out
└── prices
    ├── _SUCCESS
    └── month=2014-01-01
        ├── stock=FB
        │   └── part-00093-3741bdc2-345a-488e-82da-53bb586cd23b.c000.json
        ├── stock=GOOG
        │   └── part-00014-3741bdc2-345a-488e-82da-53bb586cd23b.c000.json
        └── stock=MSFT
            └── part-00152-3741bdc2-345a-488e-82da-53bb586cd23b.c000.json

MSFT 文件如下所示：

{"values":{"MSFT":{"1388534400":55,"1388620800":55}}}

虽然“值”列名称可能不在您的字典结构中，但我希望这说明您可以做什么。