BUG: algos.factorizes moves null values when sort=False

Author: rhshadrach

doc/source/whatsnew/v1.5.0.rst

@@ -96,6 +96,7 @@ Other enhancements
 - :meth:`pd.concat` now raises when ``levels`` contains duplicate values (:issue:`46653`)
 - Added ``numeric_only`` argument to :meth:`DataFrame.corr`, :meth:`DataFrame.corrwith`, and :meth:`DataFrame.cov` (:issue:`46560`)
 - A :class:`errors.PerformanceWarning` is now thrown when using ``string[pyarrow]`` dtype with methods that don't dispatch to ``pyarrow.compute`` methods (:issue:`42613`)
+- The method :meth:`.ExtensionArray.factorize` has gained the argument ``dropna`` for determining how null values are to be treated. (:issue:`46601`)
 
 .. ---------------------------------------------------------------------------
 .. _whatsnew_150.notable_bug_fixes:
@@ -605,6 +606,7 @@ Groupby/resample/rolling
 - Bug in :meth:`SeriesGroupBy.apply` would incorrectly name its result when there was a unique group (:issue:`46369`)
 - Bug in :meth:`.Rolling.var` would segfault calculating weighted variance when window size was larger than data size (:issue:`46760`)
 - Bug in :meth:`Grouper.__repr__` where ``dropna`` was not included. Now it is (:issue:`46754`)
+- Bug in :meth:`DataFrame.groupby` and :meth:`Series.groupby` with ``dropna=False`` and ``sort=False`` would put any null groups at the end instead the order that they are encountered (:issue:`46584`)
 
 Reshaping
 ^^^^^^^^^

pandas/_libs/hashtable_class_helper.pxi.in

@@ -658,7 +658,7 @@ cdef class {{name}}HashTable(HashTable):
                             return_inverse=return_inverse)
 
     def factorize(self, const {{dtype}}_t[:] values, Py_ssize_t na_sentinel=-1,
-                  object na_value=None, object mask=None):
+                  object na_value=None, object mask=None, ignore_na=True):
         """
         Calculate unique values and labels (no sorting!)
 
@@ -690,7 +690,7 @@ cdef class {{name}}HashTable(HashTable):
         """
         uniques_vector = {{name}}Vector()
         return self._unique(values, uniques_vector, na_sentinel=na_sentinel,
-                            na_value=na_value, ignore_na=True, mask=mask,
+                            na_value=na_value, ignore_na=ignore_na, mask=mask,
                             return_inverse=True)
 
     def get_labels(self, const {{dtype}}_t[:] values, {{name}}Vector uniques,
@@ -1037,7 +1037,7 @@ cdef class StringHashTable(HashTable):
                             return_inverse=return_inverse)
 
     def factorize(self, ndarray[object] values, Py_ssize_t na_sentinel=-1,
-                  object na_value=None, object mask=None):
+                  object na_value=None, object mask=None, ignore_na=True):
         """
         Calculate unique values and labels (no sorting!)
 
@@ -1067,7 +1067,7 @@ cdef class StringHashTable(HashTable):
         """
         uniques_vector = ObjectVector()
         return self._unique(values, uniques_vector, na_sentinel=na_sentinel,
-                            na_value=na_value, ignore_na=True,
+                            na_value=na_value, ignore_na=ignore_na,
                             return_inverse=True)
 
     def get_labels(self, ndarray[object] values, ObjectVector uniques,
@@ -1290,7 +1290,7 @@ cdef class PyObjectHashTable(HashTable):
                             return_inverse=return_inverse)
 
     def factorize(self, ndarray[object] values, Py_ssize_t na_sentinel=-1,
-                  object na_value=None, object mask=None):
+                  object na_value=None, object mask=None, ignore_na=True):
         """
         Calculate unique values and labels (no sorting!)
 
@@ -1320,7 +1320,7 @@ cdef class PyObjectHashTable(HashTable):
         """
         uniques_vector = ObjectVector()
         return self._unique(values, uniques_vector, na_sentinel=na_sentinel,
-                            na_value=na_value, ignore_na=True,
+                            na_value=na_value, ignore_na=ignore_na,
                             return_inverse=True)
 
     def get_labels(self, ndarray[object] values, ObjectVector uniques,

pandas/core/algorithms.py

@@ -502,6 +502,7 @@ def factorize_array(
     size_hint: int | None = None,
     na_value=None,
     mask: npt.NDArray[np.bool_] | None = None,
+    dropna: bool = True,
 ) -> tuple[npt.NDArray[np.intp], np.ndarray]:
     """
     Factorize a numpy array to codes and uniques.
@@ -523,6 +524,11 @@ def factorize_array(
         If not None, the mask is used as indicator for missing values
         (True = missing, False = valid) instead of `na_value` or
         condition "val != val".
+    dropna: bool, default True
+        Whether null values will appear in uniques. When False, null values
+        will receive a nonnegative code instead of na_sentinel.
+
+        ..versionadded:: 1.5.0
 
     Returns
     -------
@@ -541,7 +547,11 @@ def factorize_array(
 
     table = hash_klass(size_hint or len(values))
     uniques, codes = table.factorize(
-        values, na_sentinel=na_sentinel, na_value=na_value, mask=mask
+        values,
+        na_sentinel=na_sentinel,
+        na_value=na_value,
+        mask=mask,
+        ignore_na=dropna,
     )
 
     # re-cast e.g. i8->dt64/td64, uint8->bool
@@ -728,25 +738,35 @@ def factorize(
 
     if not isinstance(values.dtype, np.dtype):
         # i.e. ExtensionDtype
-        codes, uniques = values.factorize(na_sentinel=na_sentinel)
+        # TODO: pass ignore_na=dropna. When sort is True we ignore_na here and append
+        #       on the end because safe_sort does not handle null values in uniques
+        codes, uniques = values.factorize(
+            na_sentinel=na_sentinel, dropna=dropna or sort
+        )
     else:
         values = np.asarray(values)  # convert DTA/TDA/MultiIndex
+        # TODO: pass ignore_na=dropna; see above
         codes, uniques = factorize_array(
-            values, na_sentinel=na_sentinel, size_hint=size_hint
+            values,
+            na_sentinel=na_sentinel,
+            size_hint=size_hint,
+            dropna=dropna or sort,
         )
 
     if sort and len(uniques) > 0:
         uniques, codes = safe_sort(
             uniques, codes, na_sentinel=na_sentinel, assume_unique=True, verify=False
         )
 
-    code_is_na = codes == na_sentinel
-    if not dropna and code_is_na.any():
-        # na_value is set based on the dtype of uniques, and compat set to False is
-        # because we do not want na_value to be 0 for integers
-        na_value = na_value_for_dtype(uniques.dtype, compat=False)
-        uniques = np.append(uniques, [na_value])
-        codes = np.where(code_is_na, len(uniques) - 1, codes)
+    if not dropna and sort:
+        # TODO: Can remove if we pass ignore_na=dropna; see above
+        code_is_na = codes == na_sentinel
+        if code_is_na.any():
+            # na_value is set based on the dtype of uniques, and compat set to False is
+            # because we do not want na_value to be 0 for integers
+            na_value = na_value_for_dtype(uniques.dtype, compat=False)
+            uniques = np.append(uniques, [na_value])
+            codes = np.where(code_is_na, len(uniques) - 1, codes)
 
     uniques = _reconstruct_data(uniques, original.dtype, original)
 

pandas/core/arrays/arrow/array.py

@@ -105,8 +105,11 @@ def copy(self: ArrowExtensionArrayT) -> ArrowExtensionArrayT:
         return type(self)(self._data)
 
     @doc(ExtensionArray.factorize)
-    def factorize(self, na_sentinel: int = -1) -> tuple[np.ndarray, ExtensionArray]:
-        encoded = self._data.dictionary_encode()
+    def factorize(
+        self, na_sentinel: int = -1, dropna=True
+    ) -> tuple[np.ndarray, ExtensionArray]:
+        null_encoding = "mask" if dropna else "encode"
+        encoded = self._data.dictionary_encode(null_encoding=null_encoding)
         indices = pa.chunked_array(
             [c.indices for c in encoded.chunks], type=encoded.type.index_type
         ).to_pandas()

pandas/core/arrays/base.py

@@ -1002,14 +1002,21 @@ def _values_for_factorize(self) -> tuple[np.ndarray, Any]:
         """
         return self.astype(object), np.nan
 
-    def factorize(self, na_sentinel: int = -1) -> tuple[np.ndarray, ExtensionArray]:
+    def factorize(
+        self, na_sentinel: int = -1, dropna: bool = True
+    ) -> tuple[np.ndarray, ExtensionArray]:
         """
         Encode the extension array as an enumerated type.
 
         Parameters
         ----------
         na_sentinel : int, default -1
             Value to use in the `codes` array to indicate missing values.
+        dropna: bool, default True
+            Whether null values will appear in uniques. When False, null values
+            will receive a nonnegative code instead of na_sentinel.
+
+            ..versionadded:: 1.5.0
 
         Returns
         -------
@@ -1044,7 +1051,7 @@ def factorize(self, na_sentinel: int = -1) -> tuple[np.ndarray, ExtensionArray]:
         arr, na_value = self._values_for_factorize()
 
         codes, uniques = factorize_array(
-            arr, na_sentinel=na_sentinel, na_value=na_value
+            arr, na_sentinel=na_sentinel, na_value=na_value, dropna=dropna
         )
 
         uniques_ea = self._from_factorized(uniques, self)

pandas/core/arrays/datetimelike.py

@@ -1878,7 +1878,7 @@ def _with_freq(self, freq):
 
     # --------------------------------------------------------------
 
-    def factorize(self, na_sentinel=-1, sort: bool = False):
+    def factorize(self, na_sentinel=-1, sort: bool = False, dropna: bool = True):
         if self.freq is not None:
             # We must be unique, so can short-circuit (and retain freq)
             codes = np.arange(len(self), dtype=np.intp)
@@ -1888,7 +1888,8 @@ def factorize(self, na_sentinel=-1, sort: bool = False):
                 uniques = uniques[::-1]
             return codes, uniques
         # FIXME: shouldn't get here; we are ignoring sort
-        return super().factorize(na_sentinel=na_sentinel)
+        result = super().factorize(na_sentinel=na_sentinel, dropna=dropna)
+        return result
 
 
 # -------------------------------------------------------------------

pandas/core/arrays/masked.py

@@ -868,16 +868,38 @@ def searchsorted(
         return self._data.searchsorted(value, side=side, sorter=sorter)
 
     @doc(ExtensionArray.factorize)
-    def factorize(self, na_sentinel: int = -1) -> tuple[np.ndarray, ExtensionArray]:
+    def factorize(
+        self, na_sentinel: int = -1, dropna: bool = True
+    ) -> tuple[np.ndarray, ExtensionArray]:
         arr = self._data
         mask = self._mask
 
-        codes, uniques = factorize_array(arr, na_sentinel=na_sentinel, mask=mask)
+        codes, uniques = factorize_array(
+            arr, na_sentinel=na_sentinel, mask=mask, dropna=True
+        )
 
         # check that factorize_array correctly preserves dtype.
         assert uniques.dtype == self.dtype.numpy_dtype, (uniques.dtype, self.dtype)
 
-        uniques_ea = type(self)(uniques, np.zeros(len(uniques), dtype=bool))
+        # Make room for a null value if we're not ignoring it and it exists
+        size = len(uniques) if dropna or not mask.any() else len(uniques) + 1
+        uniques_mask = np.zeros(size, dtype=bool)
+        if not dropna:
+            na_index = mask.argmax()
+            if mask[na_index]:
+                # Insert na with the proper code
+                na_code = 0 if na_index == 0 else codes[:na_index].argmax() + 1
+                if na_sentinel < 0:
+                    # codes can never equal na_sentinel and be >= na_code
+                    codes[codes >= na_code] += 1
+                else:
+                    codes[(codes >= na_code) & (codes != na_sentinel)] += 1
+                codes[codes == na_sentinel] = na_code
+                # dummy value for uniques; not used since uniques_mask will be True
+                uniques = np.insert(uniques, na_code, 0)
+                uniques_mask[na_code] = True
+        uniques_ea = type(self)(uniques, uniques_mask)
+
         return codes, uniques_ea
 
     @doc(ExtensionArray._values_for_argsort)

pandas/core/arrays/sparse/array.py

@@ -847,13 +847,18 @@ def _values_for_factorize(self):
         # Still override this for hash_pandas_object
         return np.asarray(self), self.fill_value
 
-    def factorize(self, na_sentinel: int = -1) -> tuple[np.ndarray, SparseArray]:
+    def factorize(
+        self, na_sentinel: int = -1, dropna: bool = True
+    ) -> tuple[np.ndarray, SparseArray]:
         # Currently, ExtensionArray.factorize -> Tuple[ndarray, EA]
         # The sparsity on this is backwards from what Sparse would want. Want
         # ExtensionArray.factorize -> Tuple[EA, EA]
         # Given that we have to return a dense array of codes, why bother
         # implementing an efficient factorize?
-        codes, uniques = algos.factorize(np.asarray(self), na_sentinel=na_sentinel)
+        na_sentinel_arg = na_sentinel if dropna else None
+        codes, uniques = algos.factorize(np.asarray(self), na_sentinel=na_sentinel_arg)
+        if not dropna:
+            codes[codes == -1] = na_sentinel
         uniques_sp = SparseArray(uniques, dtype=self.dtype)
         return codes, uniques_sp
 

pandas/core/arrays/string_.py

@@ -382,6 +382,13 @@ def _values_for_factorize(self):
         arr[mask] = -1
         return arr, -1
 
+    @classmethod
+    def _from_factorized(cls, values, original):
+        assert values.dtype == original._ndarray.dtype
+        # When dropna (i.e. ignore_na) is False, can get -1 from nulls
+        values[values == -1] = None
+        return original._from_backing_data(values)
+
     def __setitem__(self, key, value):
         value = extract_array(value, extract_numpy=True)
         if isinstance(value, type(self)):

pandas/core/groupby/grouper.py

@@ -657,9 +657,10 @@ def group_index(self) -> Index:
 
     @cache_readonly
     def _codes_and_uniques(self) -> tuple[npt.NDArray[np.signedinteger], ArrayLike]:
-        if self._passed_categorical:
+        if self._dropna and self._passed_categorical:
             # we make a CategoricalIndex out of the cat grouper
-            # preserving the categories / ordered attributes
+            # preserving the categories / ordered attributes;
+            # doesn't (yet) handle dropna=False
             cat = self.grouping_vector
             categories = cat.categories
 

pandas/tests/groupby/test_groupby_dropna.py

@@ -348,3 +348,52 @@ def test_groupby_nan_included():
         tm.assert_numpy_array_equal(result_values, expected_values)
     assert np.isnan(list(result.keys())[2])
     assert list(result.keys())[0:2] == ["g1", "g2"]
+
+
+@pytest.mark.parametrize(
+    "key",
+    [
+        pd.Series([2, np.nan, 1, 2]),
+        pd.Series([2, np.nan, 1, 2], dtype="UInt8"),
+        pd.Series([2, np.nan, 1, 2], dtype="Int8"),
+        pd.Series([2, np.nan, 1, 2], dtype="UInt16"),
+        pd.Series([2, np.nan, 1, 2], dtype="Int16"),
+        pd.Series([2, np.nan, 1, 2], dtype="UInt32"),
+        pd.Series([2, np.nan, 1, 2], dtype="Int32"),
+        pd.Series([2, np.nan, 1, 2], dtype="UInt64"),
+        pd.Series([2, np.nan, 1, 2], dtype="Int64"),
+        pd.Series([2, np.nan, 1, 2], dtype="Float32"),
+        pd.Series([2, np.nan, 1, 2], dtype="Int64"),
+        pd.Series([2, np.nan, 1, 2], dtype="Float64"),
+        pd.Series(["y", None, "x", "y"], dtype="category"),
+        pd.Series(["y", pd.NA, "x", "y"], dtype="string"),
+        pd.Series(["y", pd.NA, "x", "y"], dtype="string[pyarrow]"),
+        pd.Series(
+            ["2016-01-01", np.datetime64("NaT"), "2017-01-01", "2016-01-01"],
+            dtype="datetime64[ns]",
+        ),
+        pd.Series(
+            [
+                pd.Period("2012-02-01", freq="D"),
+                pd.NA,
+                pd.Period("2012-01-01", freq="D"),
+                pd.Period("2012-02-01", freq="D"),
+            ]
+        ),
+        pd.Series(pd.arrays.SparseArray([2, np.nan, 1, 2])),
+    ],
+)
+@pytest.mark.parametrize("test_series", [True, False])
+def test_no_sort_keep_na(key, test_series):
+    df = pd.DataFrame({"key": key, "a": [1, 2, 3, 4]})
+    gb = df.groupby("key", dropna=False, sort=False)
+    if test_series:
+        gb = gb["a"]
+    result = gb.sum()
+    expected = pd.DataFrame({"a": [5, 2, 3]}, index=key[:-1].rename("key"))
+    if test_series:
+        expected = expected["a"]
+    if expected.index.is_categorical():
+        # TODO: Slicing reorders categories?
+        expected.index = expected.index.reorder_categories(["y", "x"])
+    tm.assert_equal(result, expected)

pandas/tests/groupby/transform/test_transform.py

@@ -1311,12 +1311,8 @@ def test_transform_cumcount():
 
 
 @pytest.mark.parametrize("keys", [["A1"], ["A1", "A2"]])
-def test_null_group_lambda_self(request, sort, dropna, keys):
+def test_null_group_lambda_self(sort, dropna, keys):
     # GH 17093
-    if not sort and not dropna:
-        msg = "GH#46584: null values get sorted when sort=False"
-        request.node.add_marker(pytest.mark.xfail(reason=msg, strict=False))
-
     size = 50
     nulls1 = np.random.choice([False, True], size)
     nulls2 = np.random.choice([False, True], size)