| 785 | |
| 786 | |
| 787 | def test_squeeze(): |
| 788 | df = pd.DataFrame({"x": [1, 3, 6]}) |
| 789 | df2 = pd.DataFrame({"x": [0]}) |
| 790 | s = pd.Series({"test": 0, "b": 100}) |
| 791 | |
| 792 | ddf = dd.from_pandas(df, 3) |
| 793 | ddf2 = dd.from_pandas(df2, 3) |
| 794 | ds = dd.from_pandas(s, 2) |
| 795 | |
| 796 | assert_eq(df.squeeze(), ddf.squeeze()) |
| 797 | assert_eq(pd.Series([0], name="x"), ddf2.squeeze()) |
| 798 | assert_eq(ds.squeeze(), s.squeeze()) |
| 799 | |
| 800 | with pytest.raises(NotImplementedError) as info: |
| 801 | ddf.squeeze(axis=0) |
| 802 | msg = f"{type(ddf)} does not support squeeze along axis 0" |
| 803 | assert msg in str(info.value) |
| 804 | |
| 805 | with pytest.raises(ValueError) as info: |
| 806 | ddf.squeeze(axis=2) |
| 807 | msg = f"No axis {2} for object type {type(ddf)}" |
| 808 | assert msg in str(info.value) |
| 809 | |
| 810 | with pytest.raises(ValueError) as info: |
| 811 | ddf.squeeze(axis="test") |
| 812 | msg = f"No axis test for object type {type(ddf)}" |
| 813 | assert msg in str(info.value) |
| 814 | |
| 815 | |
| 816 | def test_where_mask(): |