Apply suggestions from code review

ibis-project · Aug 2, 2024 · edc7838 · edc7838
1 parent 13973de
commit edc7838
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 16 deletions.
diff --git a/ibis/backends/pyspark/tests/test_udf.py b/ibis/backends/pyspark/tests/test_udf.py
@@ -45,14 +45,14 @@ def test_python_udf(t, df):
     tm.assert_frame_equal(result, expected)
 
 
-@pytest.mark.skipif(PYSPARK_LT_35, reason="pyarrow UDFs require PySpark 3.5+")
+@pytest.mark.xfail(PYSPARK_LT_35, reason="pyarrow UDFs require PySpark 3.5+")
 def test_pyarrow_udf(t, df):
     result = t.mutate(repeated=pyarrow_repeat(t.str_col, 2)).execute()
     expected = df.assign(repeated=df.str_col * 2)
     tm.assert_frame_equal(result, expected)
 
 
-@pytest.mark.skipif(not PYSPARK_LT_35, reason="pyarrow UDFs require PySpark 3.5+")
+@pytest.mark.xfail(not PYSPARK_LT_35, reason="pyarrow UDFs require PySpark 3.5+")
 def test_illegal_udf_type(t):
     @ibis.udf.scalar.pyarrow
     def my_add_one(x) -> str:

diff --git a/ibis/backends/tests/test_udf.py b/ibis/backends/tests/test_udf.py
@@ -6,12 +6,6 @@
 
 import ibis.common.exceptions as com
 from ibis import _, udf
-
-try:
-    from ibis.backends.pyspark import PYSPARK_LT_35
-except ImportError:
-    # Workaround when pyspark is not installed
-    PYSPARK_LT_35 = True
 from ibis.backends.tests.errors import Py4JJavaError
 
 no_python_udfs = mark.notimpl(
@@ -37,13 +31,6 @@
     reason="Docker image has Python 3.10, results in `cloudpickle` version mismatch",
 )
 
-pyspark_version_mismatch = mark.notimpl(
-    ["pyspark"],
-    condition=PYSPARK_LT_35,
-    raises=NotImplementedError,
-    reason="pyarrow UDFs require PySpark 3.5+",
-)
-
 
 @no_python_udfs
 @cloudpickle_version_mismatch
@@ -188,7 +175,7 @@ def add_one_pyarrow(s: int) -> int:  # s is series, int is the element type
                     raises=NotImplementedError,
                     reason="backend doesn't support pyarrow UDFs",
                 ),
-                pyspark_version_mismatch,
+                mark.xfail_version(pyspark=["pyspark<3.5"]),
             ],
         ),
     ],