[backport 3.0.x] BUG: use fill_null fallback for bug in pyarrow 21 on Windows (fixes join, fillna, duplicated, etc with pyarrow-backed arrays) (#64081) (#64194)

Co-authored-by: Jeongmin Gil <puduck622@gmail.com>

Author: jorisvandenbossche

doc/source/whatsnew/v3.0.1.rst

@@ -33,6 +33,7 @@ Bug fixes
 - Fixed a bug in the :func:`comparison_op` raising a ``TypeError`` for zerodim
   subclasses of ``np.ndarray`` (:issue:`63205`)
 - Added additional typing aliases in :py:mod:`pandas.api.typing.aliases` (:issue:`64098`)
+- Fixed bug in :func:`merge` where NaN values in pyarrow-backed string dtype join keys were incorrectly matched with non-NaN values on Windows with pyarrow 21 (:issue:`64060`)
 - Fixed bug in :meth:`DataFrame.loc` when setting new row and new columns simultaneously filling existing columns with ``b''`` instead of ``NaN`` (:issue:`58316`)
 - Fixed thread safety issues in :class:`DataFrame` internals on the free-threaded build (:issue:`63685`).
 - Prevent buffer overflow in :meth:`Rolling.corr` and :meth:`Rolling.cov` with variable windows when passing ``other`` with a longer index than the original window. This now raises ``ValueError`` (:issue:`62937`)

pandas/compat/pyarrow.py

@@ -2,6 +2,9 @@
 
 from __future__ import annotations
 
+import sys
+from typing import Any
+
 from pandas.util.version import Version
 
 PYARROW_MIN_VERSION = "13.0.0"
@@ -32,3 +35,57 @@
     pa_version_under21p0 = True
     pa_version_under22p0 = True
     HAS_PYARROW = False
+
+
+def _safe_fill_null(
+    arr: pa.Array | pa.ChunkedArray, fill_value: Any
+) -> pa.Array | pa.ChunkedArray:
+    """
+    Safe wrapper for pyarrow.compute.fill_null with fallback for Windows + pyarrow 21.
+
+    pyarrow 21.0.0 on Windows has a bug in fill_null that incorrectly fills null values.
+    This function uses a fallback implementation for that specific case, otherwise uses
+    the standard pyarrow.compute.fill_null.
+
+    Parameters
+    ----------
+    arr : pyarrow.Array | pyarrow.ChunkedArray
+        Input array with potential null values.
+    fill_value : Any
+        Value to fill nulls with.
+
+    Returns
+    -------
+    pyarrow.Array | pyarrow.ChunkedArray
+        Array with nulls filled with fill_value.
+    """
+    import pyarrow.compute as pc
+
+    is_windows = sys.platform in ["win32", "cygwin"]
+    use_fallback = (
+        HAS_PYARROW and is_windows and not pa_version_under21p0 and pa_version_under22p0
+    )
+    if not use_fallback or isinstance(fill_value, (pa.Array, pa.ChunkedArray)):
+        return pc.fill_null(arr, fill_value)
+
+    fill_scalar = pa.scalar(fill_value, type=arr.type)
+
+    if pa.types.is_duration(arr.type):
+
+        def fill_null_duration(arr: pa.Array, fill_scalar: pa.Scalar) -> pa.Array:
+            mask = pc.is_null(arr)
+            zero_duration = pa.scalar(0, type=arr.type)
+            arr_zeroed = pc.if_else(mask, zero_duration, arr)
+            return pc.if_else(mask, fill_scalar, arr_zeroed)
+
+        if isinstance(arr, pa.ChunkedArray):
+            return pa.chunked_array(
+                [fill_null_duration(chunk, fill_scalar) for chunk in arr.chunks]
+            )
+        return fill_null_duration(arr, fill_scalar)
+
+    if isinstance(arr, pa.ChunkedArray):
+        return pa.chunked_array(
+            [pc.if_else(pc.is_null(chunk), fill_scalar, chunk) for chunk in arr.chunks]
+        )
+    return pc.if_else(pc.is_null(arr), fill_scalar, arr)

pandas/core/arrays/arrow/array.py

@@ -95,6 +95,8 @@
     import pyarrow as pa
     import pyarrow.compute as pc
 
+    from pandas.compat.pyarrow import _safe_fill_null
+
     from pandas.core.dtypes.dtypes import ArrowDtype
 
     ARROW_CMP_FUNCS = {
@@ -1404,7 +1406,7 @@ def fillna(
 
         try:
             return self._from_pyarrow_array(
-                pc.fill_null(self._pa_array, fill_value=fill_value)
+                _safe_fill_null(self._pa_array, fill_value=fill_value)
             )
         except pa.ArrowNotImplementedError:
             # ArrowNotImplementedError: Function 'coalesce' has no kernel
@@ -1470,7 +1472,7 @@ def factorize(
             combined = encoded.combine_chunks()
             pa_indices = combined.indices
             if pa_indices.null_count > 0:
-                pa_indices = pc.fill_null(pa_indices, -1)
+                pa_indices = _safe_fill_null(pa_indices, -1)
             indices = pa_indices.to_numpy(zero_copy_only=False, writable=True).astype(
                 np.intp, copy=False
             )
@@ -1929,7 +1931,7 @@ def _str_accumulate(
                 return self._from_pyarrow_array(pa_array)
             if skipna:
                 if name == "cumsum":
-                    pa_array = pc.fill_null(pa_array, "")
+                    pa_array = _safe_fill_null(pa_array, "")
                 else:
                     # We can retain the running min/max by forward/backward filling.
                     pa_array = pc.fill_null_forward(pa_array)

pandas/core/reshape/merge.py

@@ -2837,7 +2837,8 @@ def _factorize_keys(
             isinstance(lk.dtype, StringDtype) and lk.dtype.storage == "pyarrow"
         ):
             import pyarrow as pa
-            import pyarrow.compute as pc
+
+            from pandas.compat.pyarrow import _safe_fill_null
 
             len_lk = len(lk)
             lk = lk._pa_array  # type: ignore[attr-defined]
@@ -2849,10 +2850,10 @@ def _factorize_keys(
             )
 
             llab, rlab, count = (
-                pc.fill_null(dc.indices[slice(len_lk)], -1)
+                _safe_fill_null(dc.indices[slice(len_lk)], -1)
                 .to_numpy()
                 .astype(np.intp, copy=False),
-                pc.fill_null(dc.indices[slice(len_lk, None)], -1)
+                _safe_fill_null(dc.indices[slice(len_lk, None)], -1)
                 .to_numpy()
                 .astype(np.intp, copy=False),
                 len(dc.dictionary),

scripts/validate_unwanted_patterns.py

@@ -61,6 +61,7 @@
     "_DatetimeTZBlock",
     "_check_pyarrow_available",
     "_parser",  # https://github.com/pandas-dev/pandas/issues/60833
+    "_safe_fill_null",
 }