test pandas & arrow entity queries

lenskit · Jan 21, 2025 · 321f6c1 · 321f6c1
1 parent 0661acd
commit 321f6c1
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/lenskit/tests/data/test_dataset_entity_query.py b/lenskit/tests/data/test_dataset_entity_query.py
@@ -12,6 +12,10 @@ def test_all_entities(rng: Generator, ml_ratings: pd.DataFrame, ml_ds: Dataset):
     assert np.all(ml_ds.entities("item").ids() == np.unique(ml_ratings["item_id"]))
     assert np.all(ml_ds.entities("item").numbers() == np.arange(ml_ds.item_count))
 
+    df = ml_ds.entities("item").pandas()
+    assert len(df) == ml_ratings["item_id"].nunique()
+    assert np.all(df["item_id"] == ml_ds.items.ids())
+
 
 def test_entity_subset_ids(rng: Generator, ml_ratings: pd.DataFrame, ml_ds: Dataset):
     item_ids = rng.choice(ml_ratings["item_id"].unique(), 20, replace=False)
@@ -21,6 +25,9 @@ def test_entity_subset_ids(rng: Generator, ml_ratings: pd.DataFrame, ml_ds: Data
     assert np.all(ents.ids() == item_ids)
     assert np.all(ents.numbers() == ml_ds.items.numbers(item_ids))
 
+    df = ents.pandas()
+    assert np.all(df["item_id"] == item_ids)
+
 
 def test_entity_subset_numbers(rng: Generator, ml_ratings: pd.DataFrame, ml_ds: Dataset):
     inos = rng.choice(ml_ratings["item_id"].nunique(), 20, replace=False)