lkpy/tests/data/test_builder_entities.py at 596d1f1ab4d0ca0bd3837d39450bda9df8d152e6 · lenskit/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
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
# This file is part of LensKit.
# Copyright (C) 2018-2023 Boise State University.
# Copyright (C) 2023-2025 Drexel University.
# Licensed under the MIT license, see LICENSE.md for details.
# SPDX-License-Identifier: MIT

# pyright: strict
import numpy as np
import pyarrow as pa
import pandas as pd

from pytest import raises

from lenskit.data import DatasetBuilder
from lenskit.diagnostics import DataError
from lenskit.testing import ml_test_dir


def test_empty_builder():
    dsb = DatasetBuilder()
    assert not dsb.relationship_classes()
    ecs = dsb.entity_classes()
    assert len(ecs) == 1
    items = ecs["item"]
    assert items is not None
    assert items.id_type is None
    assert not items.attributes
    assert dsb.record_count("item") == 0
    with raises(KeyError):
        dsb.record_count("user")


def test_add_entity_ids():
    dsb = DatasetBuilder()

    dsb.add_entities("item", ["a", "b", "x", "y", "z"])
    icls = dsb.entity_classes()["item"]
    assert icls.id_type == "str"
    assert dsb.entity_id_type("item") == pa.string()

    ds = dsb.build()
    assert ds.item_count == 5
    assert np.all(ds.items.ids() == ["a", "b", "x", "y", "z"])
    assert np.all(np.asarray(ds.entities("item").ids()) == ["a", "b", "x", "y", "z"])


def test_add_new_entity_ids():
    dsb = DatasetBuilder()

    dsb.add_entities("user", ["a", "b", "x", "y", "z"])
    ecs = dsb.entity_classes()
    assert set(ecs.keys()) == {"user", "item"}

    ucls = ecs["user"]
    assert ucls.id_type == "str"
    assert dsb.entity_id_type("user") == pa.string()

    ds = dsb.build()
    assert ds.item_count == 0
    assert ds.user_count == 5
    assert np.all(ds.users.ids() == ["a", "b", "x", "y", "z"])
    assert np.all(ds.entities("user").ids() == ["a", "b", "x", "y", "z"])


def test_add_invalid_entity_class_name():
    dsb = DatasetBuilder()
    with raises(ValueError, match="invalid"):
        dsb.add_entities("_item", ["a", "b", "c"])

    dsb.add_entities("item", ["a", "b", "c"])
    ds = dsb.build()
    assert ds.item_count == 3
    assert np.all(ds.items.ids() == ["a", "b", "c"])


def test_add_duplicate_entities_forbidden():
    dsb = DatasetBuilder()

    dsb.add_entities("item", ["a", "b", "c"])
    with raises(DataError, match="duplicate"):
        dsb.add_entities("item", ["d", "b", "e"])


def test_add_duplicate_entities_overwrite():
    dsb = DatasetBuilder()

    dsb.add_entities("item", ["a", "b", "c"])
    dsb.add_entities("item", ["d", "b", "e"], duplicates="overwrite")

    ds = dsb.build()
    assert ds.item_count == 5
    assert set(ds.items.ids()) == {"a", "b", "c", "d", "e"}


def test_add_entities_upcast_existing():
    dsb = DatasetBuilder()

    dsb.add_entities("item", np.arange(10, dtype="i4"))
    assert dsb.entity_id_type("item") == pa.int32()

    dsb.add_entities("item", np.arange(100, 110, dtype="i8"))
    assert dsb.entity_id_type("item") == pa.int64()

    ds = dsb.build()
    assert ds.items.ids().dtype == np.int64


def test_add_entities_upcast_new():
    dsb = DatasetBuilder()

    dsb.add_entities("item", np.arange(10, dtype="i8"))
    assert dsb.entity_id_type("item") == pa.int64()

    dsb.add_entities("item", np.arange(100, 110, dtype="i4"))
    assert dsb.entity_id_type("item") == pa.int64()

    ds = dsb.build()
    assert ds.items.ids().dtype == np.int64


def test_reject_invalid_entity_id_type():
    dsb = DatasetBuilder()

    with raises(TypeError):
        dsb.add_entities("item", np.random.randn(10))  # type: ignore


def test_reject_duplicate_ids():
    dsb = DatasetBuilder()

    with raises(DataError):
        dsb.add_entities("item", ["a", "b", "a"])


def test_add_entities_twice():
    dsb = DatasetBuilder()

    dsb.add_entities("user", ["a", "b", "x", "y", "z"])
    ecs = dsb.entity_classes()
    assert set(ecs.keys()) == {"user", "item"}

    dsb.add_entities("user", ["q", "r", "s"])

    ds = dsb.build()
    assert ds.item_count == 0
    assert ds.user_count == 8
    assert np.all(ds.users.ids() == ["a", "b", "x", "y", "z", "q", "r", "s"])


def test_add_entities_with_dataframe():
    dsb = DatasetBuilder()

    items = pd.read_csv(ml_test_dir / "movies.csv")
    items = items.rename(columns={"movieId": "item_id"}).set_index("item_id")

    genres = items["genres"].str.split("|")
    items["genres"] = genres

    dsb.add_entities("item", items)

    ds = dsb.build()

    assert ds.entities("item").attribute("title").is_scalar
    assert ds.entities("item").attribute("genres").is_list


def test_add_entities_with_arrow_table():
    dsb = DatasetBuilder()

    items = pd.read_csv(ml_test_dir / "movies.csv")
    items = items.rename(columns={"movieId": "item_id"}).set_index("item_id")

    genres = items["genres"].str.split("|")
    items["genres"] = genres
    table = pa.Table.from_pandas(items)

    dsb.add_entities("item", table)

    ds = dsb.build()

    assert ds.entities("item").attribute("title").is_scalar
    assert ds.entities("item").attribute("genres").is_list