Source code for featuretools.primitives.standard.aggregation.num_unique

import pandas as pd
from woodwork.column_schema import ColumnSchema
from woodwork.logical_types import IntegerNullable

from featuretools.primitives.base.aggregation_primitive_base import AggregationPrimitive


class NumUnique(AggregationPrimitive):
    """    确定不同值的数量,忽略 `NaN` 值.

    Args:

        use_string_for_pd_calc (bool): 确定是使用字符串 'nunique' 还是函数
            pd.Series.nunique 来进行原始计算.此参数的设置是为了
            解决 bug https://github.com/pandas-dev/pandas/issues/57317.
            默认使用字符串.

    Examples:

        >>> num_unique = NumUnique(use_string_for_pd_calc=False)
        >>> num_unique(['red', 'blue', 'green', 'yellow'])
        4

        `NaN` 值将被忽略.

        >>> num_unique(['red', 'blue', 'green', 'yellow', None])
        4"""

    name = "num_unique"
    input_types = [ColumnSchema(semantic_tags={"category"})]
    return_type = ColumnSchema(logical_type=IntegerNullable, semantic_tags={"numeric"})
    stack_on_self = False
    description_template = "the number of unique elements in {}"

    def __init__(self, use_string_for_pd_calc=True):
        self.use_string_for_pd_calc = use_string_for_pd_calc

    def get_function(self):
        if self.use_string_for_pd_calc:
            return "nunique"
        return pd.Series.nunique