我将尝试为此提供一个尴尬的数组解决方案。
首先,我们需要知道您的数据集中的哪些内容真正在扩展。我知道您可能拥有数百万个块,但这意味着您不应使用具有唯一字符串值键的 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-like”)操作到达一半,这表明该库需要一个“groupby”函数。
当然,你可以在 Numba 中完成所有事情,这往往是整体上最快的解决方案,但这个问题关键取决于分组和唯一性,在 Python 中最自然地使用 dicts 或 set 完成,但很难让 Numba推断除纯数字和数组之外的任何事物的类型。(为了 JIT 编译代码并使其快速运行,Numba 必须知道它们的类型,但 Python 类型注释足够新,它们仍然被合并到 Numba 中。)
因此,让我们从将上述数据结构变成一个 Awkward Array 开始。根据您的数据源,有不同的方法可以做到这一点(这部分文档是完整的),但我只是让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
现在我们可以利用每个单元内的笛卡尔积来制作 bc 对 ( 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”值正确连接,您就想混合一个块中的所有单元,因为您对每个块中“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已记录在案。我们最终希望将这些 bc 对按 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”功能,但根本没有。这将是一个自然的添加,也许应该是一个功能请求。
所以在这一点上,我们切换到 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
snapshot
>>> 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”,则需要使用 adepth_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()
打印 a"["
和output.begin_record()
打印 a"{"
等)
@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 的类型化 dict,但它抱怨 dict 的值是不可散列的列表。(字典值不需要是可散列的,所以我想知道那里发生了什么。)正如唯一性和排序是 Awkward Array 的前沿,键入 dicts 是 Numba 的前沿,因为在 Python 中键入的概念本身是相当新的。
我没有对这些提议的解决方案中的任何一个进行性能测试。