使用Cython查找2D NumPy数组唯一行的最快方法

Question

我有一个2D的NumPy数组，它可以是任何类型，但对于这个例子，我们可以假设它是整数。我正在寻找一种最快的方法来查找数组中所有独特的行。

我的初始策略是将每行转换为元组，并将其添加到集合中。如果集合的长度增加，则意味着找到了一个唯一的行。

我不知道如何快速将每行哈希为字节。有一个问题在这里entire array is hashed here.

我尝试过创建元组，有很多种方法，每种方法都会影响性能。这是我的函数，我展示了4种不同的变化:

版本1:

def unique_int_tuple1(ndarray[np.int64_t, ndim=2] a):
    cdef int i, len_before
    cdef int nr = a.shape[0]
    cdef int nc = a.shape[1]
    cdef set s = set()
    cdef ndarray[np.uint8_t, cast = True] idx = np.zeros(nr, dtype='bool')

    for i in range(nr):
        len_before = len(s)
        s.add(tuple(a[i]))        # THIS LINE IS CHANGED FOR ALL VERSIONS
        if len(s) > len_before:
            idx[i] = True
    return idx

版本 2：

s.add(tuple([a[i, j] for j in range(nc)]))

第三版：

vals是一个列表，其长度等于列数。

for j in range(nc):
    vals[j] = a[i, j]
    s.add(tuple(vals))

版本 4：

s.add((a[i, 0], a[i, 1], a[i, 2], a[i, 3]))

性能

a = np.random.randint(0, 8, (10**5, 4))
%timeit unique_int_tuple1(a)
125 ms ± 1.96 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit unique_int_tuple2(a)
14.5 ms ± 93.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit unique_int_tuple3(a)
11.7 ms ± 126 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit unique_int_tuple4(a)
9.59 ms ± 108 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

避免使用元组构造函数（版本4）可以获得良好的性能提升。

使用tostring：

从上面链接的SO问题中，我可以在每一行上使用tostring方法，然后进行哈希处理。

def unique_int_tostring(ndarray[np.int64_t, ndim=2] a):
    cdef int i, j
    cdef int nr = a.shape[0]
    cdef int nc = a.shape[1]
    cdef set s = set()
    cdef ndarray[np.uint8_t, cast = True] idx = np.zeros(nr, dtype='bool')

    for i in range(nr):
        len_before = len(s)
        s.add(a[i].tostring())
        if len(s) > len_before:
            idx[i] = True
    return idx

这个可以工作，但非常慢：

%timeit unique_int_tostring(a)
40 ms ± 428 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

使用类型化的memoryview

我认为，减速的一个重要原因是访问每一行 a[i]。我们可以使用类型化的memoryviews来提高性能，但我不知道如何将类型化的memoryviews元素转换为字符串，以便对它们进行哈希。

def unique_int_memoryview(long[:, :] a):
    cdef int i, j
    cdef int nr = a.shape[0]
    cdef int nc = a.shape[1]
    cdef set s = set()
    for i in range(nr):
        s.add(<SOMETHING>)   # NO IDEA HERE
    return s

- Ted Petrou

---

你可以通过使用a[i,:]而不是a[i]（对于ndarray和内存视图都适用）来获得一些改进 - 尽管我怀疑这不会有太大的改善。 - DavidW

@DavidW 不幸的是，那并没有帮助。其他想法包括在循环之前将整个数组转换为字符串。此外，我不确定将一行转换为字符串是否能保证唯一性。 - Ted Petrou

对我来说，np.unique(a, axis=0) 的输出是 20.6 毫秒 ± 77.5 微秒每次循环（平均值±7 次运行的标准差，每次 10 次循环）。也许这可以作为起点？我不确定这种方法是否可以在 Cython 中使用。 - roganjosh

@roganjosh np.unique很慢（除非数据中有很少的重复项）而且首先对数据进行排序。我正在寻找一种基于哈希的解决方案。 - Ted Petrou

为什么不使用自己的哈希函数？我已经成功地在纯Cython中实现了FNV哈希：https://github.com/yt-project/yt/blob/c1569367c6e3d8d0a02e10d0f3d0bd701d2e2114/yt/utilities/lib/fnv_hash.pyx - ngoldbaum

Answer 1

您可以使用ndarray.view()将dtype更改为byte string，然后使用pandas.Series.duplicated()查找重复行：

import numpy as np

a = np.random.randint(0, 5, size=(200, 3))
s = pd.Series(a.view(("S", a[0].nbytes))[:, 0])
s.duplicated()

< p > < code > duplicated() 的核心算法是由Cython实现的。但是它需要将原始数组转换为对象数组，这可能会很慢。

为了跳过对象数组，您可以直接使用Pandas使用的khash库，以下是该C代码：

#include "khash.h"

typedef struct _Buf{
    unsigned short n;
    char * pdata;
} Buf;

khint32_t kh_buf_hash_func(Buf key)
{
    int i;
    char * s;
    khint32_t hash = 0;
    s = key.pdata;
    for(i=0;i<key.n;i++)
    {
        hash += *s++;
        hash += (hash << 10);
        hash ^= (hash >> 6);
    }
    hash += (hash << 3);
    hash ^= (hash >> 11);
    hash += (hash << 15);    
    return hash;
}

khint32_t kh_buf_hash_equal(Buf a, Buf b)
{
    int i;
    if(a.n != b.n) return 0;
    for(i=0;i<a.n;i++){
        if(a.pdata[i] != b.pdata[i]) return 0;
    }
    return 1;
}

KHASH_INIT(buf, Buf, char, 0, kh_buf_hash_func, kh_buf_hash_equal)


void duplicated(char * arr, int row_size, int count, char * res)
{
    kh_buf_t * khbuf;
    Buf row;
    int i, absent;
    khint_t k;
    row.n = row_size;

    khbuf = kh_init_buf();
    kh_resize_buf(khbuf, 4 * count);

    for(i=0;i<count;i++){
        row.pdata = &arr[i * row_size];
        k = kh_put_buf(khbuf, row, &absent);
        if (absent){
            res[i] = 0;
        }
        else{
            res[i] = 1;
        }
    }    
    kh_destroy_buf(khbuf);
}

然后使用Cython、Ctypes或cffi对duplicated()函数进行包装。

- HYRY

---

@TedPetrou，我修改了答案，包括可以由Cython包装的c代码。 - HYRY

Answer 2

这个让我惊讶的是，这种方法速度较慢，但不管怎样，这里有一个C++解决方案，可以按字节将每行哈希为一组。 "诀窍"在于获取元素 <char*>&a[i, 0] 的地址，大部分工作都是簿记。 我可能会做一些明显的次优选择和/或使用不同的哈希表实现性能更好。

编辑：

关于如何从行创建字符串 我认为你能做的最好的方法是 - 从指针构造一个 bytes 对象。 这确实涉及到一行的复制，请参阅 C API 文档。

%%cython
from numpy cimport *
cimport numpy as np
import numpy as np
from cpython.bytes cimport PyBytes_FromStringAndSize

def unique_int_string(ndarray[np.int64_t, ndim=2] a):
    cdef int i, len_before
    cdef int nr = a.shape[0]
    cdef int nc = a.shape[1]
    cdef set s = set()
    cdef ndarray[np.uint8_t, cast = True] idx = np.zeros(nr, dtype='bool')
    cdef bytes string

    for i in range(nr):
        len_before = len(s)
        string = PyBytes_FromStringAndSize(<char*>&a[i, 0], sizeof(np.int64_t) * nc)
        s.add(string)
        if len(s) > len_before:
            idx[i] = True
    return idx

// 时间控制

In [9]: from unique import unique_ints

In [10]: %timeit unique_int_tuple4(a)
100 loops, best of 3: 10.1 ms per loop

In [11]: %timeit unique_ints(a)
100 loops, best of 3: 11.9 ms per loop

In [12]: (unique_ints(a) == unique_int_tuple4(a)).all()
Out[12]: True

// 助手.h

#include <unordered_set>
#include <cstring>

struct Hasher {
    size_t size;
    size_t operator()(char* buf) const {
        // https://github.com/yt-project/yt/blob/c1569367c6e3d8d0a02e10d0f3d0bd701d2e2114/yt/utilities/lib/fnv_hash.pyx
        size_t hash_val = 2166136261;
        for (int i = 0; i < size; ++i) {
                hash_val ^= buf[i];
                hash_val *= 16777619;
        }
        return hash_val;
    }
};
struct Comparer {
    size_t size;
    bool operator()(char* lhs, char* rhs) const {
        return (std::memcmp(lhs, rhs, size) == 0) ? true : false;
    }
};

struct ArraySet {
    std::unordered_set<char*, Hasher, Comparer> set;

    ArraySet (size_t size) : set(0, Hasher{size}, Comparer{size}) {}
    ArraySet () {}

    bool add(char* buf) {
        auto p = set.insert(buf);
        return p.second;
    }
};

// unique.pyx

from numpy cimport int64_t, uint8_t
import numpy as np

cdef extern from 'helper.h' nogil:
    cdef cppclass ArraySet:
        ArraySet()
        ArraySet(size_t)
        bint add(char*)


def unique_ints(int64_t[:, :] a):
    cdef:
        Py_ssize_t i, nr = a.shape[0], nc = a.shape[1]
        ArraySet s = ArraySet(sizeof(int64_t) * nc)
        uint8_t[:] idx = np.zeros(nr, dtype='uint8')

        bint found;

    for i in range(nr):
        found = s.add(<char*>&a[i, 0])
        if found:
            idx[i] = True

    return idx

// 设置.py

from setuptools import setup, Extension
from Cython.Build import cythonize
import numpy as np

exts = [
  Extension('unique', ['unique.pyx'], language='c++', include_dirs=[np.get_include()])
]

setup(name='test', ext_modules=cythonize(exts))

- chrisb

---

std::hash 对于向量默认情况下也未定义，因此在这种情况下必须定义自定义哈希函数。 - chrisb

哇，你真是个巫师。非常酷。我测试了新的解决方案，它比整数元组慢了50%。难道没有更快的方法将行转换为字符串吗？@HYRY上面的解决方案使用a.view(("S", a[0].nbytes))。 - Ted Petrou

我认为这已经是最快的了（可能有误！）- Python字符串拥有自己的内存，因此它们必须复制缓冲区，而@HYRY正在创建一个numpy字节类型，该类型可以查看原始数据，类似于我尝试使用C++版本。 - chrisb

好的，这太棒了。我对时间的估计错了。在我的实际数据集中，使用您的版本可以获得更好的性能，因此唯一值的数量和行长度必须会有所不同。对于浮点数，性能也要好得多（在一个大部分唯一的10**5乘以10的数组上，性能提高了4倍）。 - Ted Petrou

顺便说一下，我正在尝试构建一个更简单、更高效的 pandas 版本，名为 dexplo，并且已经使所有操作变得更快，除了按浮点数和大范围整数进行分组。通过迭代集合，字符串分组已经更快了，现在这个新功能应该会让它更加领先。如果您有兴趣帮忙，请告诉我。 - Ted Petrou