python - PySpark:从数据帧创建字典的字典？

Question

我有以下格式的数据，这些数据是从 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 在 native Python 中工作，但由于数据量巨大，我宁愿在集群/Spark 上完成此操作。

Answer 1

在 Spark 2.4 中，您可以在聚合每只股票的值时使用 map_from_arrays 构建日期值映射。然后只需使用 create_map 将股票代码用作键即可。此示例使用 python 3.4 中的 ChainMap 来构建您所描述的最终字典结构。

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
    }
}

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

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

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}}}

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