lkpy/tests/accel/test_scatter.py at 89dbc65d95f89f7a4df0bb781b135b79491f1163 · mdekstrand/lkpy · 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
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
# This file is part of LensKit.
# Copyright (C) 2018-2023 Boise State University.
# Copyright (C) 2023-2026 Drexel University.
# Licensed under the MIT license, see LICENSE.md for details.
# SPDX-License-Identifier: MIT

from typing import Any

import numpy as np
import pyarrow as pa

import hypothesis.extra.numpy as nph
import hypothesis.strategies as st
from hypothesis import given
from pytest import raises

from lenskit._accel import data


@given(
    st.data(),
    nph.arrays(
        st.one_of(
            nph.integer_dtypes(endianness="="),
            nph.unsigned_integer_dtypes(endianness="="),
            nph.floating_dtypes(endianness="=", sizes=(16, 32, 64)),
        ),
        nph.array_shapes(max_dims=1),
    ),
    st.sampled_from([np.int32, np.int64]),
)
def test_scatter_dst_array(hd: st.DataObject, dst: np.ndarray[tuple[int], Any], idx_t: np.dtype):
    size = len(dst)
    idx = np.asarray(
        list(hd.draw(st.sets(st.integers(min_value=0, max_value=size - 1)))), dtype=idx_t
    )
    src = hd.draw(nph.arrays(dst.dtype, len(idx)))

    dst_a = pa.array(dst)
    idx_a = pa.array(idx)
    src_a = pa.array(src)

    arr_a = data.scatter_array(dst_a, idx_a, src_a)
    assert isinstance(arr_a, pa.Array)
    assert arr_a.type == dst_a.type
    assert arr_a.null_count == 0

    arr = arr_a.to_numpy()

    assert np.array_equal(arr[idx], src, equal_nan=True)


@given(
    st.data(),
    st.integers(0, 16 * 1024 + 1),
    st.sampled_from([np.int32, np.int64]),
)
def test_scatter_dst_size(hd: st.DataObject, size, idx_t: np.dtype):
    if size:
        idx = np.asarray(
            list(hd.draw(st.sets(st.integers(min_value=0, max_value=size - 1)))), dtype=idx_t
        )
    else:
        idx = np.asarray([], dtype=idx_t)
    src = hd.draw(
        nph.arrays(
            st.one_of(
                nph.integer_dtypes(endianness="="),
                nph.floating_dtypes(endianness="=", sizes=(16, 32, 64)),
            ),
            len(idx),
        )
    )

    idx_a = pa.array(idx)
    src_a = pa.array(src)

    arr_a = data.scatter_array_empty(size, idx_a, src_a)
    assert isinstance(arr_a, pa.Array)
    assert arr_a.type == src_a.type
    assert arr_a.null_count == size - len(src)

    arr = arr_a.to_numpy(zero_copy_only=False)

    assert np.array_equal(arr[idx], src, equal_nan=True)


def test_scatter_rejects_strings():
    strings = pa.array(["a", "b", "c", "x", "9", "0", "3", "@"])
    with raises(TypeError):
        data.scatter_array(strings, pa.array([2, 7, 0]), strings)


def test_scatter_empty_rejects_strings():
    strings = pa.array(["a", "b", "c", "x", "9", "0", "3", "@"])
    with raises(TypeError):
        data.scatter_array_empty(100, pa.array([2, 7, 0]), strings)