Skip to content
112 changes: 112 additions & 0 deletions python/cudf/cudf/core/column/numerical.py
Original file line number Diff line number Diff line change
Expand Up @@ -947,6 +947,102 @@ def nan_count(self) -> int:
return super().nan_count
return self.isnan().sum()

def isin(self, values: Sequence | ColumnBase) -> ColumnBase:
if isinstance(self.dtype, (np.dtype, pd.ArrowDtype)):
return super().isin(values)
# Mirror pandas' BaseMaskedArray.isin for masked (nullable
# integer/float/boolean) dtypes:
# * matching is done on the underlying numpy values, so e.g. a
# boolean element equals the integer 1,
# * an NA element is considered present only when pd.NA itself
# is one of the passed values (a plain NaN/None/NaT does not
# match), and
# * the result is a nullable BooleanDtype with no missing values.
# Validity-only nulls (ColumnBase.isnull): a genuine NaN value in a
# masked float column is data, not NA - it must keep matching a NaN
# needle instead of being folded into the mask
# (NumericalColumn.isnull would count it as null).
na_mask = ColumnBase.isnull(self)
# View the column through its numpy dtype (identical physical
# layout). astype is not used because casting a masked column with
# nulls to numpy raises in pandas-compatible mode, and its
# short-circuit mutates the live column's dtype in place; fillna is
# not used because it folds genuine NaN values into the fill. Null
# rows come back False from ``isin`` (the needles below never
# contain nulls) and are then corrected from ``na_mask``.
data_col = ColumnBase.create(self.plc_column, self.dtype.numpy_dtype)
cleaned_values, values_have_na = self._process_values_for_masked_isin(
values
)
if len(cleaned_values) == 0:
# Nothing can match; also avoids the object-dtype column an
# empty needle list produces, which an all-null data column
# cannot be compared against in pandas-compatible mode.
result = as_column(
False, length=len(self), dtype=np.dtype(np.bool_)
)
else:
result = data_col.isin(cleaned_values)
if na_mask.any():
result = result | na_mask if values_have_na else result & ~na_mask
return result.astype(pd.BooleanDtype())

@staticmethod
def _process_values_for_masked_isin(
values: Sequence | ColumnBase,
) -> tuple[Sequence, bool]:
"""Normalize ``values`` for a masked ``isin``.

Drops NA-like needles (they are handled through the mask by the
caller) and reports whether ``pd.NA`` itself was among them.
"""
# The normalization below rebinds across container types, so use a
# loosely typed alias.
host_values: Any = values
# Bring device-backed ``values`` to host in a single transfer; the
# per-element inspection below would otherwise read them
# element-wise or fail on non-iterable cudf objects.
if isinstance(host_values, (ColumnBase, cudf.Series, cudf.Index)):
host_values = host_values.to_pandas()
elif isinstance(host_values, cp.ndarray):
host_values = cp.asnumpy(host_values)
elif isinstance(host_values, (pa.Array, pa.ChunkedArray)):
host_values = host_values.to_pylist()
if isinstance(
host_values,
(pd.Series, pd.Index, pd.api.extensions.ExtensionArray),
):
# np.asarray mirrors pandas' BaseMaskedArray.isin: a masked (or
# other extension) container decays to numpy, where pd.NA
# becomes NaN and no longer matches NA rows - only object
# containers can carry pd.NA identity.
host_values = np.asarray(host_values)
elif not isinstance(host_values, np.ndarray):
# Materialize one-shot iterators so the single pass below is
# the only traversal.
host_values = list(host_values)
# NA-like sentinels never match a real value and would otherwise
# raise when mixed with numeric values, so they are dropped from
# the needles; NA rows are handled by the caller through the mask
# instead.
if isinstance(host_values, np.ndarray) and host_values.dtype != object:
# A non-object array cannot hold pd.NA, and NaT is the only
# NA-like value it can hold (NaN is a matchable value), so no
# Python-level scan is needed.
if host_values.dtype.kind in "mM":
host_values = host_values[~np.isnat(host_values)]
return cast("Sequence", host_values), False
# ``host_values`` is a list or an object array here; a single
# Python pass mirrors pandas' pd.NA identity scan.
cleaned: list[Any] = []
values_have_na = False
for value in host_values:
if value is pd.NA:
values_have_na = True
elif not is_na_like(value):
cleaned.append(value)
return cleaned, values_have_na

def _process_values_for_isin(
self, values: Sequence | ColumnBase
) -> tuple[ColumnBase, ColumnBase]:
Expand All @@ -964,6 +1060,22 @@ def _process_values_for_isin(
rhs = rhs.astype(lhs.dtype)
elif lhs.can_cast_safely(rhs.dtype):
lhs = lhs.astype(rhs.dtype)
elif (
isinstance(lhs.dtype, np.dtype)
and isinstance(rhs.dtype, np.dtype)
and lhs.dtype.kind in "biuf"
and rhs.dtype.kind in "biuf"
and "b" in (lhs.dtype.kind, rhs.dtype.kind)
):
# A boolean column compares by value against numeric needles
# (``True == 1``) like numpy/pandas. ``can_cast_safely`` reports
# bool<->numeric as unsafe, so promote the boolean side to the
# numeric dtype (a bool always fits) rather than bailing out to
# an all-False result.
if lhs.dtype.kind == "b":
lhs = lhs.astype(rhs.dtype)
else:
rhs = rhs.astype(lhs.dtype)
return lhs, rhs

def _can_return_nan(self, skipna: bool | None = None) -> bool:
Expand Down
12 changes: 0 additions & 12 deletions python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py
Original file line number Diff line number Diff line change
Expand Up @@ -3802,18 +3802,6 @@ def pytest_unconfigure(config):
"tests/series/methods/test_info.py::test_info_series[False]": "TODO: Add a reason for failure",
"tests/series/methods/test_info.py::test_info_series[True]": "TODO: Add a reason for failure",
"tests/series/methods/test_info.py::test_info_show_counts_false": "assert '<class 'pand...: 140.0 bytes' == '<class 'pand...ge: 8.0 bytes'",
"tests/series/methods/test_isin.py::TestSeriesIsIn::test_isin_masked_types[data0-values0-expected0-Float64]": "TODO: Add a reason for failure",
"tests/series/methods/test_isin.py::TestSeriesIsIn::test_isin_masked_types[data0-values0-expected0-Int64]": "TODO: Add a reason for failure",
"tests/series/methods/test_isin.py::TestSeriesIsIn::test_isin_masked_types[data0-values0-expected0-boolean]": "TODO: Add a reason for failure",
"tests/series/methods/test_isin.py::TestSeriesIsIn::test_isin_masked_types[data2-values2-expected2-Float64]": "TODO: Add a reason for failure",
"tests/series/methods/test_isin.py::TestSeriesIsIn::test_isin_masked_types[data2-values2-expected2-Int64]": "TODO: Add a reason for failure",
"tests/series/methods/test_isin.py::TestSeriesIsIn::test_isin_masked_types[data2-values2-expected2-boolean]": "TODO: Add a reason for failure",
"tests/series/methods/test_isin.py::TestSeriesIsIn::test_isin_masked_types[data4-values4-expected4-Float64]": "TODO: Add a reason for failure",
"tests/series/methods/test_isin.py::TestSeriesIsIn::test_isin_masked_types[data4-values4-expected4-Int64]": "TODO: Add a reason for failure",
"tests/series/methods/test_isin.py::TestSeriesIsIn::test_isin_masked_types[data4-values4-expected4-boolean]": "TODO: Add a reason for failure",
"tests/series/methods/test_isin.py::TestSeriesIsIn::test_isin_masked_types[data5-values5-expected5-Int64]": "TODO: Add a reason for failure",
"tests/series/methods/test_isin.py::TestSeriesIsIn::test_isin_masked_types[data5-values5-expected5-boolean]": "TODO: Add a reason for failure",
"tests/series/methods/test_isin.py::test_isin_large_series_and_pdNA[boolean-data3-values3-expected3]": "AssertionError: Attributes of Series are different",
"tests/series/methods/test_map.py::test_map_callable[MockEngineDecorator]": "AssertionError: assert Index([], dtype='object', name='bar') is Index([], dtype='object', name='bar')",
"tests/series/methods/test_map.py::test_map_callable[None]": "AssertionError: assert Index([], dtype='object', name='bar') is Index([], dtype='object', name='bar')",
"tests/series/methods/test_map.py::test_map_categorical_na_action[None-expected0]": "TODO: Add a reason for failure",
Expand Down
185 changes: 185 additions & 0 deletions python/cudf/cudf/tests/series/methods/test_isin.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

import cupy as cp
import numpy as np
import pandas as pd
import pyarrow as pa
Expand Down Expand Up @@ -162,6 +163,190 @@ def test_isin_categorical(data, values):
assert_eq(got, expected)


@pytest.mark.parametrize("dtype", ["boolean", "Int64", "Float64"])
@pytest.mark.parametrize(
"data,values",
[
([0, 1, 0], [1]),
([0, 1, 0], [1, pd.NA]),
([0, pd.NA, 0], [1, 0]),
([0, 1, pd.NA], [1, pd.NA]),
([0, 1, pd.NA], [1, np.nan]),
([0, pd.NA, pd.NA], [np.nan, pd.NaT, None]),
],
)
def test_isin_masked_types(dtype, data, values):
# Series.isin on a pandas masked (nullable integer/float/boolean) dtype
# returns a nullable BooleanDtype result and matches pandas' NA semantics:
# * comparison is done on the underlying values (a boolean element equals
# the integer 1), and
# * an NA element is considered present only when pd.NA itself is one of
# ``values`` (a plain NaN/None/NaT does not match).
psr = pd.Series(data, dtype=dtype)
gsr = cudf.Series(data, dtype=dtype)

got = gsr.isin(values)
expected = psr.isin(values)

assert got.dtype == pd.BooleanDtype()
assert_eq(got, expected)


@pytest.mark.parametrize(
"values",
[[1], [0], [1, 0], [2], [1.0], [1.5], [True], [False]],
)
def test_isin_bool_against_numeric(values):
# A boolean Series compares equal to the integers/floats 0 and 1, matching
# numpy/pandas value semantics (previously cudf returned all-False for a
# numeric ``values`` argument).
psr = pd.Series([True, False, False, True])
gsr = cudf.Series([True, False, False, True])

got = gsr.isin(values)
assert got.dtype == np.dtype("bool")
assert_eq(got, psr.isin(values))


@pytest.mark.parametrize(
"values",
[
lambda: (value for value in [1, 2]),
lambda: cudf.Series([1, 2]),
lambda: cudf.Index([1, 2]),
lambda: cp.array([1, 2]),
lambda: pa.chunked_array([[1, 2]]),
],
ids=["generator", "cudf-series", "cudf-index", "cupy", "pyarrow"],
)
def test_isin_masked_values_containers(values):
# ``values`` is materialized on host once for masked dtypes:
# device-backed inputs are moved in a single transfer instead of
# element-wise reads (cudf objects are not Python-iterable at all).
psr = pd.Series([0, 1, pd.NA], dtype="Int64")
gsr = cudf.Series([0, 1, pd.NA], dtype="Int64")

assert_eq(gsr.isin(values()), psr.isin([1, 2]))


@pytest.mark.parametrize(
"values",
[
pd.Series([1, pd.NA], dtype="Int64"),
pd.Index([1, pd.NA], dtype="Int64"),
pd.array([1, pd.NA], dtype="Int64"),
],
ids=["pd-series", "pd-index", "pd-array"],
)
def test_isin_masked_values_container_NA_does_not_match(values):
# pandas materializes pandas-container ``values`` with np.asarray, where
# pd.NA decays to NaN: an NA needle inside a masked container does NOT
# match NA rows, unlike a literal pd.NA in a list.
psr = pd.Series([0, 1, pd.NA], dtype="Int64")
gsr = cudf.Series([0, 1, pd.NA], dtype="Int64")

assert_eq(gsr.isin(values), psr.isin(values))


def test_isin_masked_one_shot_iterator_with_NA():
# A one-shot iterator must not be exhausted before the pd.NA
# inspection: the NA row still matches when pd.NA is one of ``values``.
psr = pd.Series([0, 1, pd.NA], dtype="Int64")
gsr = cudf.Series([0, 1, pd.NA], dtype="Int64")

got = gsr.isin(value for value in [1, pd.NA])
expected = psr.isin([1, pd.NA])
assert_eq(got, expected)


def test_isin_masked_nan_is_value_not_na():
# A genuine NaN value in a masked float column is data, not NA: it
# matches a NaN needle and does not match pd.NA (mirrors pandas, where
# only the mask is NA).
psr = pd.Series(
pd.arrays.FloatingArray(
np.array([np.nan, 1.0, 0.0]),
np.array([False, False, True]),
)
)
gsr = cudf.Series([np.nan, 1.0, None], dtype="Float64", nan_as_null=False)

assert_eq(gsr.isin([np.nan]), psr.isin([np.nan]))
assert_eq(gsr.isin([pd.NA]), psr.isin([pd.NA]))


def test_isin_masked_pandas_compatible_mode():
# The masked path must work in pandas-compatible mode (the
# masked-with-nulls to numpy astype guard must not trigger) and must
# not mutate the input Series' dtype via the in-place astype
# short-circuit.
with cudf.option_context("mode.pandas_compatible", True):
gsr = cudf.Series([1, pd.NA], dtype="Int64")
got = gsr.isin([1])
assert gsr.dtype == pd.Int64Dtype()
assert got.dtype == pd.BooleanDtype()
assert got.to_pandas().tolist() == [True, False]

gsr = cudf.Series([True, pd.NA], dtype="boolean")
got = gsr.isin([True])
assert gsr.dtype == pd.BooleanDtype()
assert got.to_pandas().tolist() == [True, False]

gsr = cudf.Series([1.5, 2.5], dtype="Float64")
got = gsr.isin([1.5])
assert gsr.dtype == pd.Float64Dtype()
assert got.to_pandas().tolist() == [True, False]


@pytest.mark.parametrize("values", [[], [pd.NA], [None], [1]])
def test_isin_masked_all_na(values):
# An all-NA masked Series with needles that clean to empty used to
# raise in pandas-compatible mode (an empty needle list produces an
# object-dtype column that an all-null column cannot be cast to there).
psr = pd.Series([pd.NA, pd.NA], dtype="Int64")
expected = psr.isin(values)

gsr = cudf.Series([pd.NA, pd.NA], dtype="Int64")
assert_eq(gsr.isin(values), expected)

with cudf.option_context("mode.pandas_compatible", True):
gsr = cudf.Series([pd.NA, pd.NA], dtype="Int64")
got = gsr.isin(values)
assert_eq(got, expected)


def test_isin_masked_does_not_mutate_dtype():
# Regression test: the astype in the masked path used to flip the
# input's dtype from Int64 to int64 in pandas-compatible mode.
gsr = cudf.Series([1, 2], dtype="Int64")
gsr.isin([1])
assert gsr.dtype == pd.Int64Dtype()

with cudf.option_context("mode.pandas_compatible", True):
gsr = cudf.Series([1, 2], dtype="Int64")
gsr.isin([1])
assert gsr.dtype == pd.Int64Dtype()


@pytest.mark.parametrize(
"psr",
[
pd.Series([0, 1, 0], dtype=pd.ArrowDtype(pa.int64())),
pd.Series([1.0, 2.0, None], dtype=pd.ArrowDtype(pa.float64())),
pd.Series([True, False, True], dtype=pd.ArrowDtype(pa.bool_())),
pd.Series(["a", "b", "a"], dtype="category"),
],
)
def test_isin_non_masked_extension_returns_numpy_bool(psr):
# Arrow and categorical inputs yield a numpy bool result (only masked
# numeric/boolean dtypes upgrade to nullable boolean).
gsr = cudf.from_pandas(psr)

got = gsr.isin([psr.iloc[0]])
assert got.dtype == np.dtype("bool")
assert_eq(got, psr.isin([psr.iloc[0]]))


@pytest.mark.parametrize(
"dtype",
[
Expand Down
Loading