numstore/test.py at main · rowles/numstore · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
import numpy as np
import numstore
import time


def test_small():
    arr = np.array([111, 55, 11], dtype=np.uint64)
    arr = np.array([11, 55, 142], dtype=np.uint64)

    for algo in numstore.ALGO_STR.keys():
        w = numstore.Writer(algo, 'data/small.vec')
        w.write(arr)
        w.close()

        r = numstore.Reader('data/small.vec')
        arr2 = r.read()
        r.close()

        np.testing.assert_array_equal(arr, arr2)


def test_big():
    N = 50_000
    arr = np.array(range(N), dtype=np.uint64)

    for algo in numstore.ALGO_STR.keys():
        w = numstore.Writer(algo, 'data/big.vec')
        w.write(arr)
        w.close()

        r = numstore.Reader('data/big.vec')
        arr2 = r.read()
        r.close()

        np.testing.assert_array_equal(arr, arr2)


def test_context():
    arr = np.array([1, 2, 3], dtype=np.uint64)

    with numstore.Writer('TurboPFor', 'data/small2.vec') as w:
        w.write(arr)

    with numstore.Reader('data/small2.vec') as r:
        arr2 = r.read()

    np.testing.assert_array_equal(arr, arr2)