Improve tests for to_numeric (#10804)

hendrikmakait · web-flow · commit ca42dcf5fcd4 · 2024-01-16T12:08:39.000Z
diff --git a/dask/dataframe/tests/test_numeric.py b/dask/dataframe/tests/test_numeric.py
@@ -5,10 +5,12 @@
 import pytest
 
 from dask.array import Array, from_array
-from dask.dataframe import Series, from_pandas, to_numeric
+from dask.dataframe import Series, _dask_expr_enabled, from_pandas, to_numeric
 from dask.dataframe.utils import pyarrow_strings_enabled
 from dask.delayed import Delayed
 
+DASK_EXPR_ENABLED = _dask_expr_enabled()
+
 
 @pytest.mark.parametrize("arg", ["5", 5, "5 "])
 def test_to_numeric_on_scalars(arg):
@@ -21,6 +23,18 @@ def test_to_numeric_on_dask_array():
     arg = from_array(["1.0", "2", "-3", "5.1"])
     expected = np.array([1.0, 2.0, -3.0, 5.1])
     output = to_numeric(arg)
+    expected_dtype = "int64"
+    assert output.dtype == expected_dtype
+    assert isinstance(output, Array)
+    assert list(output.compute()) == list(expected)
+
+
+def test_to_numeric_on_dask_array_with_meta():
+    arg = from_array(["1.0", "2", "-3", "5.1"])
+    expected = np.array([1.0, 2.0, -3.0, 5.1])
+    output = to_numeric(arg, meta=np.array((), dtype="float64"))
+    expected_dtype = "float64"
+    assert output.dtype == expected_dtype
     assert isinstance(output, Array)
     assert list(output.compute()) == list(expected)
 
@@ -31,7 +45,7 @@ def test_to_numeric_on_dask_dataframe_series():
     expected = pd.to_numeric(s)
     output = to_numeric(arg)
     expected_dtype = "int64"
-    if pyarrow_strings_enabled():
+    if pyarrow_strings_enabled() and not DASK_EXPR_ENABLED:
         # `to_numeric` output depends on input dtype
         expected_dtype = "Int64"
     assert output.dtype == expected_dtype
