[Re-opened elsewhere] Allow where() to work with a Series and other=cudf.NA

python/cudf/cudf/core/_internals/where.py

@@ -3,8 +3,10 @@
 import warnings
 from typing import Any, Optional, Tuple, Union, cast
 
+import cupy
 import numpy as np
 import pandas as pd
+from numba import cuda
 
 import cudf
 from cudf._typing import ColumnLike, ScalarLike
@@ -27,7 +29,9 @@ def _normalize_scalars(col: ColumnBase, other: ScalarLike) -> ScalarLike:
             f"{type(other).__name__} to {col.dtype.name}"
         )
 
-    return cudf.Scalar(other, dtype=col.dtype if other is None else None)
+    return cudf.Scalar(
+        other, dtype=col.dtype if other in {None, cudf.NA} else None
+    )
 
 
 def _check_and_cast_columns_with_other(
@@ -234,9 +238,19 @@ def where(
 
     if isinstance(frame, DataFrame):
         if hasattr(cond, "__cuda_array_interface__"):
-            cond = DataFrame(
-                cond, columns=frame._column_names, index=frame.index
-            )
+            if (
+                isinstance(cond, DataFrame)
+                or isinstance(cond, cupy._core.core.ndarray)
+                or isinstance(cond, cuda.cudadrv.devicearray.DeviceNDArray)
+            ):
+                cond = DataFrame(
+                    cond, columns=frame._column_names, index=frame.index
+                )
+            else:
+                cond = DataFrame(
+                    {name: cond for name in frame._column_names},
+                    index=frame.index,
+                )
         elif (
             hasattr(cond, "__array_interface__")
             and cond.__array_interface__["shape"] != frame.shape

python/cudf/cudf/tests/test_dataframe.py

@@ -8736,3 +8736,23 @@ def test_frame_series_where():
     expected = gdf.where(gdf.notna(), gdf.mean())
     actual = pdf.where(pdf.notna(), pdf.mean(), axis=1)
     assert_eq(expected, actual)
+
+
+@pytest.mark.parametrize(
+    "data", [{"a": [1, 2, 3], "b": [1, 1, 0]}],
+)
+def test_frame_series_where_other(data):
+    gdf = cudf.DataFrame(data)
+    pdf = gdf.to_pandas()
+
+    expected = gdf.where(gdf["b"] == 1, cudf.NA)
+    actual = pdf.where(pdf["b"] == 1, pd.NA)
+    assert_eq(
+        actual.fillna(-1).values,
+        expected.fillna(-1).values,
+        check_dtype=False,
+    )
+
+    expected = gdf.where(gdf["b"] == 1, 0)
+    actual = pdf.where(pdf["b"] == 1, 0)
+    assert_eq(expected, actual)