From 1a96a56c1b9eebeeb97e8c22b992288b3fe66d7f Mon Sep 17 00:00:00 2001 From: Patrick Hoefler <61934744+phofl@users.noreply.github.com> Date: Thu, 8 Feb 2024 12:49:49 +0100 Subject: [PATCH] Adjust test for dask-expr --- dask/dataframe/tests/test_dataframe.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py index 82d63ff4182..06e3abdcfb8 100644 --- a/dask/dataframe/tests/test_dataframe.py +++ b/dask/dataframe/tests/test_dataframe.py @@ -1828,7 +1828,9 @@ def test_assign(): # divisions unknown won't work with pandas with pytest.raises(ValueError): - ddf_unknown.assign(c=df.a + 1) + q = ddf_unknown.assign(c=df.a + 1) + if DASK_EXPR_ENABLED: + q.optimize() # unsupported type with pytest.raises(TypeError): @@ -1836,10 +1838,14 @@ def test_assign(): # Fails when assigning known divisions to unknown divisions with pytest.raises(ValueError): - ddf_unknown.assign(foo=ddf.a) + q = ddf_unknown.assign(foo=ddf.a) + if DASK_EXPR_ENABLED: + q.optimize() # Fails when assigning unknown divisions to known divisions with pytest.raises(ValueError): - ddf.assign(foo=ddf_unknown.a) + q = ddf.assign(foo=ddf_unknown.a) + if DASK_EXPR_ENABLED: + q.optimize() df = pd.DataFrame({"A": [1, 2]}) df.assign(B=lambda df: df["A"], C=lambda df: df.A + df.B)