使用 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)]))

版本 3：

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)

避免使用 tuple 构造函数（第 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)

使用类型化的内存视图

我相信，减速的很大一部分是对每一行 a[i] 的访问。我们可以使用类型化内存视图来提高性能，但是我不知道如何将类型化内存视图的元素转换为字符串以便它们可以被散列。

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

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()

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

跳过object array，可以直接使用Pandas使用的khash library，这里是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() 函数。

Answer 2

令我惊讶的是，这速度较慢，但​​不管它值多少钱，这里有一个 c++ 解决方案，它可以完成您所指的操作 - 将每一行哈希为一组字节。 “诀窍”是获取元素<char*>&a[i, 0] 的地址——其他大部分内容都是簿记。

我可能正在做一些明显的次优和/或使用不同的哈希表实现可能会更好。

编辑：

re: 如何从一行创建一个字符串 我认为你能做的最好的就是这个 - 从指针构造一个 bytes 对象。这必然涉及行的副本，请参阅 c api docs。

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

// helper.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

// setup.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))