[GitHub] [spark] ueshin commented on a change in pull request #33474: [SPARK-36249][PYTHON] Add remove_categories to CategoricalAccessor and CategoricalIndex

[GitBox <gi...@apache.org>]

ueshin commented on a change in pull request #33474:
URL: https://github.com/apache/spark/pull/33474#discussion_r674514929



##########
File path: python/pyspark/pandas/categorical.py
##########
@@ -348,8 +348,96 @@ def as_unordered(self, inplace: bool = False) -> Optional["ps.Series"]:
         """
         return self._set_ordered(ordered=False, inplace=inplace)
 
-    def remove_categories(self, removals: pd.Index, inplace: bool = False) -> "ps.Series":
-        raise NotImplementedError()
+    def remove_categories(
+        self, removals: Union[pd.Index, Any, List], inplace: bool = False
+    ) -> Optional["ps.Series"]:
+        """
+        Remove the specified categories.
+
+        `removals` must be included in the old categories. Values which were in
+        the removed categories will be set to NaN
+
+        Parameters
+        ----------
+        removals : category or list of categories
+           The categories which should be removed.
+        inplace : bool, default False
+           Whether or not to remove the categories inplace or return a copy of
+           this categorical with removed categories.
+
+        Returns
+        -------
+        Series or None
+            Categorical with removed categories or None if ``inplace=True``.
+
+        Raises
+        ------
+        ValueError
+            If the removals are not contained in the categories
+
+        Examples
+        --------
+        >>> s = ps.Series(list("abbccc"), dtype="category")
+        >>> s  # doctest: +SKIP
+        0    a
+        1    b
+        2    b
+        3    c
+        4    c
+        5    c
+        dtype: category
+        Categories (3, object): ['a', 'b', 'c']
+
+        >>> s.cat.remove_categories('b')  # doctest: +SKIP
+        0      a
+        1    NaN
+        2    NaN
+        3      c
+        4      c
+        5      c
+        dtype: category
+        Categories (2, object): ['a', 'c']
+        """
+        if is_list_like(removals):
+            categories = [cat for cat in removals if cat is not None]  # type: List
+        elif removals is None:
+            categories = []
+        else:
+            categories = [removals]
+
+        if any(cat not in self.categories for cat in categories):
+            raise ValueError(
+                "removals must all be in old categories: {{{cats}}}".format(
+                    cats=", ".join(
+                        set(str(cat) for cat in categories if cat not in self.categories)
+                    )
+                )
+            )
+
+        if len(categories) == 0:
+            if inplace:
+                return None
+            else:
+                psser = self._data
+                return psser._with_new_scol(
+                    psser.spark.column, field=psser._internal.data_fields[0]
+                )

Review comment:
       `psser.copy()` keeps the anchor DataFrame as-is, but we don't need here, so explicitly drop the extra columns.




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org