Fix pytests due to loc bug

Signed-off-by: Ayush Dattagupta <[email protected]>
NVIDIA · Apr 23, 2024 · 9f38aa1 · 9f38aa1
1 parent e467078
commit 9f38aa1
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/tests/test_filters.py b/tests/test_filters.py
@@ -149,7 +149,9 @@ def test_retain_score_filter(self, letter_count_data):
         filtered_data = filter_step(letter_count_data)
 
         expected_indices = [2, 3]
-        expected_data = DocumentDataset(letter_count_data.df.loc[expected_indices])
+        # Compute before loc due to https://github.com/dask/dask-expr/issues/1036
+        expected_data = letter_count_data.df.compute().loc[expected_indices]
+        expected_data = DocumentDataset(dd.from_pandas(expected_data, 2))
         expected_data.df[score_field] = pd.Series([5, 7], index=expected_data.df.index)
         assert all_equal(
             expected_data, filtered_data
@@ -168,7 +170,9 @@ def test_filter(self, letter_count_data):
         filtered_data = filter_step(scored_data)
 
         expected_indices = [2, 3]
-        expected_data = letter_count_data.df.loc[expected_indices]
+        # Compute before loc due to https://github.com/dask/dask-expr/issues/1036
+        expected_data = letter_count_data.df.compute().loc[expected_indices]
+        expected_data = dd.from_pandas(expected_data, 2)
         expected_data[score_field] = pd.Series([5, 7], index=expected_data.index)
         expected_data = DocumentDataset(expected_data)
         assert all_equal(