我将尝试为此提供一个尴尬的数组解决方案。
首先,我们需要知道您的数据集中的哪些内容真正在扩大规模。我知道您可能拥有数百万个块,但这意味着您不应使用具有唯一字符串值键的 dict 来表示它们。创建字符串、比较字符串和按字符串查找内容的成本(尽管 dict 的哈希表对此有很大帮助)并不是一种很好的扩展方式,尤其是当这些字符串中的唯一信息是排序时。
在我看来,您的单位数量也是可变的(事实上,最后一个被命名为“unit10k”、“unit12k”和“unit5k”)。
最后,我重新阅读了您的问题陈述,但似乎名为“a”的字段没有进入问题。我会忽略它。
因此,我会说您的数据结构会更好:
as_python = [
# first block
[
# unit 1
{"b": 23, "c": [10]},
# unit 2
{"b": 15, "c": [20, 10]},
# ...
# unit 10k
{"b": 10, "c": [6, 10, 20, 5]},
],
# second block
[
# unit 1
{"b": 10, "c": [17, 12, 9]},
# unit 2
{"b": 15, "c": [17, 15, 9, 4, 12]},
# ...
# unit 12k
{"b": 23, "c": [12, 22, 15, 4]},
],
# ...
# millionth block
[
# unit 1
{"b": 64, "c": [50]},
# unit 2
{"b": 55, "c": [89, 59, 77]},
# ...
# unit 15k
{"b": 85, "c": [48, 90, 27, 59]},
],
]
对此的 Pythonic 解决方案是
results = []
for block in as_python:
block_result = {}
for unit in block:
b = unit["b"] # only look up b once
for c in unit["c"]:
if c not in block_result:
block_result[c] = []
block_result[c].append(b)
results.append(block_result)
导致
[
{10: [23, 15, 10], 20: [15, 10], 6: [10], 5: [10]},
{17: [10, 15], 12: [10, 15, 23], 9: [10, 15], 15: [15, 23], 4: [15, 23], 22: [23]},
{50: [64], 89: [55], 59: [55, 85], 77: [55], 48: [85], 90: [85], 27: [85]},
]
应该很快。 (仅使用内置 Python 类型的短循环对于未编译的动态类型虚拟机来说速度惊人。根据我的经验,仅使用内置 Python 类型是关键。)
至于尴尬的数组,在调用 Numba(JIT 编译的for 循环)之前,我只能使用矢量化(即“类 NumPy”)操作到达一半,这表明该库需要一个“groupby”函数。
当然,你可以在 Numba 中完成所有事情,这往往是整体上最快的解决方案,但这个问题主要取决于分组和唯一性,在 Python 中最自然地使用字典或集合来完成,但很难做到让 Numba 推断出纯数字和数组之外的任何类型。 (为了 JIT 编译代码并使其更快,Numba 必须知道它们的类型,但 Python 类型注释足够新,它们仍然被合并到 Numba 中。)
所以让我们先把上面的数据结构变成一个尴尬的数组。根据您的数据源,有不同的方法可以做到这一点(this part of the documentation 已完成),但我会让ak.Array 遍历上述 Python 数据结构。
>>> import awkward as ak
>>> as_awkward = ak.Array(as_python)
通常,如果我们遇到“X 和 Y 的所有组合”问题,我们希望使用ak.cartesian。我们不能在这里马上使用它,因为“b”数据和“c”数据的深度不同:
>>> as_awkward = ak.Array(as_python)
>>> as_awkward.b
<Array [[23, 15, 10], ... 23], [64, 55, 85]] type='3 * var * int64'>
>>> as_awkward.c
<Array [[[10], [20, 10], ... [48, 90, 27, 59]]] type='3 * var * var * int64'>
>>> as_awkward.type
3 * var * {"b": int64, "c": var * int64}
>>> as_awkward.b.type
3 * var * int64
>>> as_awkward.c.type
3 * var * var * int64
为了使它们匹配,我们可以使用np.newaxis(又名None,但我喜欢明确)创建一个新轴:
>>> import numpy as np
>>> bprime = as_awkward.b[:, :, np.newaxis]
>>> bprime
<Array [[[23], [15], [10, ... 64], [55], [85]]] type='3 * var * 1 * int64'>
>>> bprime.type
3 * var * 1 * int64
>>> as_awkward.c.type
3 * var * var * int64
现在我们可以利用每个单元内的笛卡尔积来制作 b-c 对列表 (axis=2)。
>>> combinations = ak.cartesian({"c": as_awkward.c, "b": bprime}, axis=2)
>>> combinations
<Array [[[{c: 10, b: 23}], ... c: 59, b: 85}]]] type='3 * var * var * {"c": int6...'>
>>> combinations.type
3 * var * var * {"c": int64, "b": int64}
>>> combinations.tolist()
[
[
[
[{"c": 10, "b": 23}]
],
[
[{"c": 20, "b": 15}],
[{"c": 10, "b": 15}]],
[
[{"c": 6, "b": 10}],
[{"c": 10, "b": 10}],
[{"c": 20, "b": 10}],
[{"c": 5, "b": 10}],
],
],
[
[
[{"c": 17, "b": 10}],
[{"c": 12, "b": 10}],
[{"c": 9, "b": 10}]
],
[
[{"c": 17, "b": 15}],
[{"c": 15, "b": 15}],
[{"c": 9, "b": 15}],
[{"c": 4, "b": 15}],
[{"c": 12, "b": 15}],
],
[
[{"c": 12, "b": 23}],
[{"c": 22, "b": 23}],
[{"c": 15, "b": 23}],
[{"c": 4, "b": 23}],
],
],
[
[
[{"c": 50, "b": 64}]
],
[
[{"c": 89, "b": 55}],
[{"c": 59, "b": 55}],
[{"c": 77, "b": 55}]
],
[
[{"c": 48, "b": 85}],
[{"c": 90, "b": 85}],
[{"c": 27, "b": 85}],
[{"c": 59, "b": 85}],
],
],
]
现在我们的结构太多了:只要“b”和“c”值正确连接,您就希望将所有单元混合在一个块中,因为您对每个 阻止。
>>> flattened = ak.flatten(combinations, axis=-1)
>>> flattened
<Array [[{c: 10, b: 23}, ... c: 59, b: 85}]] type='3 * var * {"c": int64, "b": i...'>
>>> flattened.type
3 * var * {"c": int64, "b": int64}
>>> flattened.tolist()
[
[
{"c": 10, "b": 23},
{"c": 20, "b": 15},
{"c": 10, "b": 15},
{"c": 6, "b": 10},
{"c": 10, "b": 10},
{"c": 20, "b": 10},
{"c": 5, "b": 10},
],
[
{"c": 17, "b": 10},
{"c": 12, "b": 10},
{"c": 9, "b": 10},
{"c": 17, "b": 15},
{"c": 15, "b": 15},
{"c": 9, "b": 15},
{"c": 4, "b": 15},
{"c": 12, "b": 15},
{"c": 12, "b": 23},
{"c": 22, "b": 23},
{"c": 15, "b": 23},
{"c": 4, "b": 23},
],
[
{"c": 50, "b": 64},
{"c": 89, "b": 55},
{"c": 59, "b": 55},
{"c": 77, "b": 55},
{"c": 48, "b": 85},
{"c": 90, "b": 85},
{"c": 27, "b": 85},
{"c": 59, "b": 85},
],
]
排序和唯一性是 Awkward Array 功能的最前沿;有一些内部编写的函数只是没有在 Python 中公开,更不用说记录了。幸运的是,ak.sort 和 ak.argsort 已记录在案。我们最终希望将这些 b-c 对按 c 分组,我们至少可以对它们进行排序:
>>> sorted = flattened[ak.argsort(flattened.c, axis=-1)]
>>> sorted
<Array [[{c: 5, b: 10}, ... {c: 90, b: 85}]] type='3 * var * {"c": int64, "b": i...'>
>>> sorted.type
3 * var * {"c": int64, "b": int64}
>>> sorted.tolist()
[
[
{"c": 5, "b": 10},
{"c": 6, "b": 10},
{"c": 10, "b": 23},
{"c": 10, "b": 15},
{"c": 10, "b": 10},
{"c": 20, "b": 15},
{"c": 20, "b": 10},
],
[
{"c": 4, "b": 15},
{"c": 4, "b": 23},
{"c": 9, "b": 10},
{"c": 9, "b": 15},
{"c": 12, "b": 10},
{"c": 12, "b": 15},
{"c": 12, "b": 23},
{"c": 15, "b": 15},
{"c": 15, "b": 23},
{"c": 17, "b": 10},
{"c": 17, "b": 15},
{"c": 22, "b": 23},
],
[
{"c": 27, "b": 85},
{"c": 48, "b": 85},
{"c": 50, "b": 64},
{"c": 59, "b": 55},
{"c": 59, "b": 85},
{"c": 77, "b": 55},
{"c": 89, "b": 55},
{"c": 90, "b": 85},
],
]
现在我们真的,真的希望我们有一个“groupby”函数,但是没有。这将是一个自然的添加,也许应该是feature request。
所以在这一点上,我们切换到 Numba。这个问题比原始问题简单得多,因为要分组的数据已经排序:我们只需要查看值何时更改以插入边界。尴尬的数组可以作为参数传递给 Numba,但它们是不可变的。要创建新数组,请使用ak.ArrayBuilder。注意:在开发这样的函数时,请分小步完成并删除 @nb.njit 以在不进行 JIT 编译的情况下进行尝试(以确保在尝试解决类型错误之前您所做的事情是正确的)。
import numba as nb
@nb.njit
def groupby(input, output):
for block in input:
output.begin_list()
output.begin_list()
last = block[0].c # note: assumes that len(block) >= 1
for unit in block:
if unit.c != last:
output.end_list()
output.begin_list()
output.append(unit)
last = unit.c
output.end_list()
output.end_list()
return output
input 是我们刚刚创建的数组 (sorted),output 是一个 ArrayBuilder,它通过 snapshot(在 Numba 之外)变成一个真正的数组。
>>> grouped = groupby(sorted, ak.ArrayBuilder()).snapshot()
>>> grouped
<Array [[[{c: 5, b: 10}], ... {c: 90, b: 85}]]] type='3 * var * var * {"c": int6...'>
>>> grouped.type
3 * var * var * {"c": int64, "b": int64}
>>> grouped.tolist()
[
[
[{"c": 5, "b": 10}],
[{"c": 6, "b": 10}],
[{"c": 10, "b": 23}, {"c": 10, "b": 15}, {"c": 10, "b": 10}],
[{"c": 20, "b": 15}, {"c": 20, "b": 10}],
],
[
[{"c": 4, "b": 15}, {"c": 4, "b": 23}],
[{"c": 9, "b": 10}, {"c": 9, "b": 15}],
[{"c": 12, "b": 10}, {"c": 12, "b": 15}, {"c": 12, "b": 23}],
[{"c": 15, "b": 15}, {"c": 15, "b": 23}],
[{"c": 17, "b": 10}, {"c": 17, "b": 15}],
[{"c": 22, "b": 23}],
],
[
[{"c": 27, "b": 85}],
[{"c": 48, "b": 85}],
[{"c": 50, "b": 64}],
[{"c": 59, "b": 55}, {"c": 59, "b": 85}],
[{"c": 77, "b": 55}],
[{"c": 89, "b": 55}],
[{"c": 90, "b": 85}],
],
]
然后您可以调整该输出以获得所需的结构。如果您想为每个列表“b”设置一个标量“c”,您需要使用depth_limit 来阻止ak.zip 将其广播到最深层次。
>>> ak.zip({"c": grouped.c[:, :, 0], "b": grouped.b}, depth_limit=2).tolist()
[
[
{"c": 5, "b": [10]},
{"c": 6, "b": [10]},
{"c": 10, "b": [23, 15, 10]},
{"c": 20, "b": [15, 10]},
],
[
{"c": 4, "b": [15, 23]},
{"c": 9, "b": [10, 15]},
{"c": 12, "b": [10, 15, 23]},
{"c": 15, "b": [15, 23]},
{"c": 17, "b": [10, 15]},
{"c": 22, "b": [23]},
],
[
{"c": 27, "b": [85]},
{"c": 48, "b": [85]},
{"c": 50, "b": [64]},
{"c": 59, "b": [55, 85]},
{"c": 77, "b": [55]},
{"c": 89, "b": [55]},
{"c": 90, "b": [85]},
],
]
如果您想直接在 Numba 中构建该类型的报告,您可以这样做。 (这只是为了说明“begin_list”等在做什么——就像“打印”一个结构,想象output.begin_list()打印一个"["和output.begin_record()打印一个"{"等)
@nb.njit
def groupby(input, output):
for block in input:
output.begin_list()
output.begin_record()
output.field("c").integer(block[0].c)
output.field("b")
output.begin_list()
last = block[0].c
for unit in block:
if unit.c != last:
output.end_list()
output.end_record()
output.begin_record()
output.field("c").integer(unit.c)
output.field("b")
output.begin_list()
output.integer(unit.b)
last = unit.c
output.end_list()
output.end_record()
output.end_list()
return output
和
>>> grouped = groupby(sorted, ak.ArrayBuilder()).snapshot()
>>> grouped
<Array [[{c: 5, b: [10]}, ... c: 90, b: [85]}]] type='3 * var * {"c": int64, "b"...'>
>>> grouped.type
3 * var * {"c": int64, "b": var * int64}
>>> grouped.tolist()
# it's the same
正如我上面所说,绝对最快的解决方案可能是在 Numba 中执行所有操作。笛卡尔积、展平和排序都创建了部分新的数组(尽可能多地重用输入,这就是为什么 Awkward 数组必须是不可变的),这涉及到数据的分配和多次传递。但是在 Numba 中很难表达涉及字典、列表和集合的问题,因为它需要类型化的字典、列表和集合。我尝试使用Numba's typed dict,但它抱怨 dict 的值是列表,它们是不可散列的。 (字典 values 不需要是可散列的,所以我想知道那里发生了什么。)正如唯一性和排序是 Awkward Array 的前沿,键入 dicts 是 Numba 概念以来的前沿用 Python 打字本身就很新。
我没有对这些建议的解决方案中的任何一个进行性能测试。