Question

我在PySpark中具有以下DataFrame df。

import pyspark.sql.functions as func

df = spark\
        .read \
        .format("org.elasticsearch.spark.sql") \
        .load("my_index/my_mapping") \
        .groupBy(["id", "type"]) \
        .agg(
            func.count(func.lit(1)).alias("number_occurrences"),
            func.countDistinct("host_id").alias("number_hosts")
        )

ds = df.collect()

我使用collect是因为分组和聚合之后的数据量总是很小并且适合内存。另外，我需要使用collect，因为我将ds传递为udf函数的参数。函数collect返回一个数组。我如何对该数组进行以下查询：对于给定的id和type，返回number_occurrences和number_hosts 。

例如，假设df包含以下行：

id   type   number_occurrences   number_hosts
1    xxx    11                   3
2    yyy    10                   4

完成df.collect()之后，如何检索等于{{1}的number_occurences和等于{{1的number_hosts的{{1}}和id }}。预期结果是：

更新：

也许有更优雅的解决方案？

type

Answer 1

如果您的id是唯一的，那么id就是这种情况，您可以根据id对数组进行排序。这只是确保顺序正确，如果您的ID是连续的，则可以直接访问记录并将ID减去1

test_df = spark.createDataFrame([
(1,"xxx",11,3),(2,"yyyy",10,4),

], ("id","type","number_occurrences","number_hosts"))
id = 1
type = "xxx"
sorted_list = sorted(test_df.collect(), cmp=lambda x,y: cmp(x["id"],y["id"]))
sorted_list[id-1]["number_occurrences"],sorted_list[id-1]["number_hosts"]

结果：

(11, 3)

如何从df.collect（）中检索特定值在PySpark中？

1 个答案: