#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements. See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
from typing import Any, Callable, List, Optional, Union, TYPE_CHECKING, cast
import pandas as pd
from pandas.api.types import (
CategoricalDtype,
is_dict_like,
is_list_like,
)
from pyspark.pandas.internal import InternalField
from pyspark.pandas.data_type_ops.categorical_ops import _to_cat
from pyspark.sql import functions as F
from pyspark.sql.types import StructField
if TYPE_CHECKING:
import pyspark.pandas as ps
class CategoricalAccessor:
"""
Accessor object for categorical properties of the Series values.
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.categories
Index(['a', 'b', 'c'], dtype='object')
>>> s.cat.codes
0 0
1 1
2 1
3 2
4 2
5 2
dtype: int8
"""
def __init__(self, series: "ps.Series"):
if not isinstance(series.dtype, CategoricalDtype):
raise ValueError("Cannot call CategoricalAccessor on type {}".format(series.dtype))
self._data = series
@property
def _dtype(self) -> CategoricalDtype:
return cast(CategoricalDtype, self._data.dtype)
@property
def categories(self) -> pd.Index:
"""
The categories of this categorical.
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.categories
Index(['a', 'b', 'c'], dtype='object')
"""
return self._dtype.categories
@categories.setter
def categories(self, categories: Union[pd.Index, List]) -> None:
dtype = CategoricalDtype(categories, ordered=self.ordered)
if len(self.categories) != len(dtype.categories):
raise ValueError(
"new categories need to have the same number of items as the old categories!"
)
internal = self._data._psdf._internal.with_new_spark_column(
self._data._column_label,
self._data.spark.column,
field=self._data._internal.data_fields[0].copy(dtype=dtype),
)
self._data._psdf._update_internal_frame(internal)
@property
def ordered(self) -> bool:
"""
Whether the categories have an ordered relationship.
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.ordered
False
"""
return self._dtype.ordered
@property
def codes(self) -> "ps.Series":
"""
Return Series of codes as well as the index.
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.codes
0 0
1 1
2 1
3 2
4 2
5 2
dtype: int8
"""
return self._data._with_new_scol(
self._data.spark.column,
field=InternalField.from_struct_field(
StructField(
name=self._data._internal.data_spark_column_names[0],
dataType=self._data.spark.data_type,
nullable=self._data.spark.nullable,
)
),
).rename()
def add_categories(self, new_categories: Union[pd.Index, Any, List]) -> Optional["ps.Series"]:
"""
Add new categories.
`new_categories` will be included at the last/highest place in the
categories and will be unused directly after this call.
Parameters
----------
new_categories : category or list-like of category
The new categories to be included.
Returns
-------
Series
Categorical with new categories added
Raises
------
ValueError
If the new categories include old categories or do not validate as
categories
See Also
--------
rename_categories : Rename categories.
reorder_categories : Reorder categories.
remove_categories : Remove the specified categories.
remove_unused_categories : Remove categories which are not used.
set_categories : Set the categories to the specified ones.
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.add_categories('x') # doctest: +SKIP
0 a
1 b
2 b
3 c
4 c
5 c
dtype: category
Categories (4, object): ['a', 'b', 'c', 'x']
"""
from pyspark.pandas.frame import DataFrame
categories: List[Any]
if is_list_like(new_categories):
categories = list(new_categories)
else:
categories = [new_categories]
if any(cat in self.categories for cat in categories):
raise ValueError(
"new categories must not include old categories: {{{cats}}}".format(
cats=", ".join(set(str(cat) for cat in categories if cat in self.categories))
)
)
internal = self._data._psdf._internal.with_new_spark_column(
self._data._column_label,
self._data.spark.column,
field=self._data._internal.data_fields[0].copy(
dtype=CategoricalDtype(list(self.categories) + categories, ordered=self.ordered)
),
)
return DataFrame(internal)._psser_for(self._data._column_label).copy()
def _set_ordered(self, *, ordered: bool) -> Optional["ps.Series"]:
from pyspark.pandas.frame import DataFrame
if self.ordered == ordered:
return self._data.copy()
else:
internal = self._data._psdf._internal.with_new_spark_column(
self._data._column_label,
self._data.spark.column,
field=self._data._internal.data_fields[0].copy(
dtype=CategoricalDtype(categories=self.categories, ordered=ordered)
),
)
return DataFrame(internal)._psser_for(self._data._column_label).copy()
def as_ordered(self) -> Optional["ps.Series"]:
"""
Set the Categorical to be ordered.
Returns
-------
Series
Ordered Categorical
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.as_ordered() # doctest: +SKIP
0 a
1 b
2 b
3 c
4 c
5 c
dtype: category
Categories (3, object): ['a' < 'b' < 'c']
"""
return self._set_ordered(ordered=True)
def as_unordered(self) -> Optional["ps.Series"]:
"""
Set the Categorical to be unordered.
Returns
-------
Series
Unordered Categorical
Examples
--------
>>> s = ps.Series(list("abbccc"), dtype="category").cat.as_ordered()
>>> 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.as_unordered() # doctest: +SKIP
0 a
1 b
2 b
3 c
4 c
5 c
dtype: category
Categories (3, object): ['a', 'b', 'c']
"""
return self._set_ordered(ordered=False)
def remove_categories(self, removals: Union[pd.Index, Any, List]) -> 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.
Returns
-------
Series
Categorical with removed categories
Raises
------
ValueError
If the removals are not contained in the categories
See Also
--------
rename_categories : Rename categories.
reorder_categories : Reorder categories.
add_categories : Add new categories.
remove_unused_categories : Remove categories which are not used.
set_categories : Set the categories to the specified ones.
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']
"""
categories: List[Any]
if is_list_like(removals):
categories = [cat for cat in removals if cat is not None]
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:
return self._data.copy()
else:
data = [cat for cat in self.categories.sort_values() if cat not in categories]
if len(data) == 0:
# We should keep original dtype when even removing all categories.
data = pd.Index(data, dtype=self.categories.dtype) # type: ignore[assignment]
dtype = CategoricalDtype(
categories=data,
ordered=self.ordered,
)
return self._data.astype(dtype)
def remove_unused_categories(self) -> Optional["ps.Series"]:
"""
Remove categories which are not used.
Returns
-------
cat : Series
Categorical with unused categories dropped
See Also
--------
rename_categories : Rename categories.
reorder_categories : Reorder categories.
add_categories : Add new categories.
remove_categories : Remove the specified categories.
set_categories : Set the categories to the specified ones.
Examples
--------
>>> s = ps.Series(pd.Categorical(list("abbccc"), categories=['a', 'b', 'c', 'd']))
>>> s # doctest: +SKIP
0 a
1 b
2 b
3 c
4 c
5 c
dtype: category
Categories (4, object): ['a', 'b', 'c', 'd']
>>> s.cat.remove_unused_categories() # doctest: +SKIP
0 a
1 b
2 b
3 c
4 c
5 c
dtype: category
Categories (3, object): ['a', 'b', 'c']
"""
categories = set(self._data.drop_duplicates()._to_pandas())
removals = [cat for cat in self.categories if cat not in categories]
categories = [cat for cat in removals if cat is not None] # type: ignore[assignment]
if len(categories) == 0:
return self._data.copy()
else:
dtype = CategoricalDtype(
[cat for cat in self.categories if cat not in categories], ordered=self.ordered
)
return self._data.astype(dtype)
def rename_categories(
self, new_categories: Union[list, dict, Callable]
) -> Optional["ps.Series"]:
"""
Rename categories.
Parameters
----------
new_categories : list-like, dict-like or callable
New categories which will replace old categories.
* list-like: all items must be unique and the number of items in
the new categories must match the existing number of categories.
* dict-like: specifies a mapping from
old categories to new. Categories not contained in the mapping
are passed through and extra categories in the mapping are
ignored.
* callable : a callable that is called on all items in the old
categories and whose return values comprise the new categories.
Returns
-------
cat : Series
Categorical with removed categories
Raises
------
ValueError
If new categories are list-like and do not have the same number of
items than the current categories or do not validate as categories
See Also
--------
reorder_categories : Reorder categories.
add_categories : Add new categories.
remove_categories : Remove the specified categories.
remove_unused_categories : Remove categories which are not used.
set_categories : Set the categories to the specified ones.
Examples
--------
>>> s = ps.Series(["a", "a", "b"], dtype="category")
>>> s.cat.rename_categories([0, 1]) # doctest: +SKIP
0 0
1 0
2 1
dtype: category
Categories (2, int64): [0, 1]
For dict-like ``new_categories``, extra keys are ignored and
categories not in the dictionary are passed through
>>> s.cat.rename_categories({'a': 'A', 'c': 'C'}) # doctest: +SKIP
0 A
1 A
2 b
dtype: category
Categories (2, object): ['A', 'b']
You may also provide a callable to create the new categories
>>> s.cat.rename_categories(lambda x: x.upper()) # doctest: +SKIP
0 A
1 A
2 B
dtype: category
Categories (2, object): ['A', 'B']
"""
from pyspark.pandas.frame import DataFrame
if is_dict_like(new_categories):
categories = [cast(dict, new_categories).get(item, item) for item in self.categories]
elif callable(new_categories):
categories = [new_categories(item) for item in self.categories]
elif is_list_like(new_categories):
if len(self.categories) != len(new_categories):
raise ValueError(
"new categories need to have the same number of items as the old categories!"
)
categories = cast(list, new_categories)
else:
raise TypeError("new_categories must be list-like, dict-like or callable.")
internal = self._data._psdf._internal.with_new_spark_column(
self._data._column_label,
self._data.spark.column,
field=self._data._internal.data_fields[0].copy(
dtype=CategoricalDtype(categories=categories, ordered=self.ordered)
),
)
return DataFrame(internal)._psser_for(self._data._column_label).copy()
def reorder_categories(
self,
new_categories: Union[pd.Index, List],
ordered: Optional[bool] = None,
) -> Optional["ps.Series"]:
"""
Reorder categories as specified in new_categories.
`new_categories` needs to include all old categories and no new category
items.
Parameters
----------
new_categories : Index-like
The categories in new order.
ordered : bool, optional
Whether or not the categorical is treated as an ordered categorical.
If not given, do not change the ordered information.
Returns
-------
cat : Series
Categorical with removed categories
Raises
------
ValueError
If the new categories do not contain all old category items or any
new ones
See Also
--------
rename_categories : Rename categories.
add_categories : Add new categories.
remove_categories : Remove the specified categories.
remove_unused_categories : Remove categories which are not used.
set_categories : Set the categories to the specified ones.
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.reorder_categories(['c', 'b', 'a'], ordered=True) # doctest: +SKIP
0 a
1 b
2 b
3 c
4 c
5 c
dtype: category
Categories (3, object): ['c' < 'b' < 'a']
"""
if not is_list_like(new_categories):
raise TypeError(
"Parameter 'new_categories' must be list-like, was '{}'".format(new_categories)
)
elif len(set(new_categories)) != len(set(self.categories)) or any(
cat not in self.categories for cat in new_categories
):
raise ValueError("items in new_categories are not the same as in old categories")
if ordered is None:
ordered = self.ordered
if new_categories == list(self.categories) and ordered == self.ordered:
return self._data.copy()
else:
dtype = CategoricalDtype(categories=new_categories, ordered=ordered)
return _to_cat(self._data).astype(dtype)
def set_categories(
self,
new_categories: Union[pd.Index, List],
ordered: Optional[bool] = None,
rename: bool = False,
) -> Optional["ps.Series"]:
"""
Set the categories to the specified new_categories.
`new_categories` can include new categories (which will result in
unused categories) or remove old categories (which results in values
set to NaN). If `rename==True`, the categories will simply be renamed
(less or more items than in old categories will result in values set to
NaN or in unused categories respectively).
This method can be used to perform more than one action of adding,
removing, and reordering simultaneously and is therefore faster than
performing the individual steps via the more specialised methods.
On the other hand this methods does not do checks (e.g., whether the
old categories are included in the new categories on a reorder), which
can result in surprising changes, for example when using special string
dtypes, which does not consider a S1 string equal to a single char
python string.
Parameters
----------
new_categories : Index-like
The categories in new order.
ordered : bool, default False
Whether or not the categorical is treated as an ordered categorical.
If not given, do not change the ordered information.
rename : bool, default False
Whether or not the new_categories should be considered as a rename
of the old categories or as reordered categories.
Returns
-------
Series with reordered categories
Raises
------
ValueError
If new_categories does not validate as categories
See Also
--------
rename_categories : Rename categories.
reorder_categories : Reorder categories.
add_categories : Add new categories.
remove_categories : Remove the specified categories.
remove_unused_categories : Remove categories which are not used.
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.set_categories(['b', 'c']) # doctest: +SKIP
0 NaN
1 b
2 b
3 c
4 c
5 c
dtype: category
Categories (2, object): ['b', 'c']
>>> s.cat.set_categories([1, 2, 3], rename=True) # doctest: +SKIP
0 1
1 2
2 2
3 3
4 3
5 3
dtype: category
Categories (3, int64): [1, 2, 3]
>>> s.cat.set_categories([1, 2, 3], rename=True, ordered=True) # doctest: +SKIP
0 1
1 2
2 2
3 3
4 3
5 3
dtype: category
Categories (3, int64): [1 < 2 < 3]
"""
from pyspark.pandas.frame import DataFrame
if not is_list_like(new_categories):
raise TypeError(
"Parameter 'new_categories' must be list-like, was '{}'".format(new_categories)
)
if ordered is None:
ordered = self.ordered
new_dtype = CategoricalDtype(new_categories, ordered=ordered)
scol = self._data.spark.column
if rename:
new_scol = (
F.when(scol >= len(new_categories), F.lit(-1).cast(self._data.spark.data_type))
.otherwise(scol)
.alias(self._data._internal.data_spark_column_names[0])
)
internal = self._data._psdf._internal.with_new_spark_column(
self._data._column_label,
new_scol,
field=self._data._internal.data_fields[0].copy(dtype=new_dtype),
)
return DataFrame(internal)._psser_for(self._data._column_label).copy()
else:
return self._data.astype(new_dtype)
def _test() -> None:
import os
import doctest
import sys
from pyspark.sql import SparkSession
import pyspark.pandas.categorical
os.chdir(os.environ["SPARK_HOME"])
globs = pyspark.pandas.categorical.__dict__.copy()
globs["ps"] = pyspark.pandas
spark = (
SparkSession.builder.master("local[4]")
.appName("pyspark.pandas.categorical tests")
.getOrCreate()
)
failure_count, test_count = doctest.testmod(
pyspark.pandas.categorical,
globs=globs,
optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE,
)
spark.stop()
if failure_count:
sys.exit(-1)
if __name__ == "__main__":
_test()
