如何在xarray Dataset中按shipID分组统计唯一品类数量

问题背景 #

现有一个名为shipDS的xarray货运项目数据集：

<xarray.Dataset>
Dimensions:   (lineItem: 10)
Dimensions without coordinates: lineItem
Data variables:
    shipID    (lineItem) int64 1 1 2 3 4 4 5 5 5 5
    prodID    (lineItem) int64 90 92 92 90 90 91 92 93 94 95
    category  (lineItem) <U1 'A' 'B' 'B' 'A' 'A' 'A' 'B' 'C' 'D' 'D'

需求：统计每个货运单（即每个shipID）包含的唯一品类数量。

创建数据集代码 #

import xarray as xr

# 构建货运数据集
# 每个货运单可能包含多个产品
shipIDs = [1,1,2,3,4,4,5,5,5,5]
productIDs = [90,92,92,90,90,91,92,93,94,95]
category = ["A","B","B","A","A","A","B","C","D","D"]

shipDS = xr.Dataset(data_vars = 
    {
        "shipID": ("lineItem", shipIDs),
        "prodID": ("lineItem", productIDs),
        "category": ("lineItem", category)
     }
)

已掌握的Pandas实现方式 #

# pandas实现代码
(
    shipDS
    .to_pandas()
    .groupby(["shipID","category"])
    .count()
    .reset_index()
    .groupby("shipID")
    .agg(numUniqCategories=('category', 'count'))
)

运行后得到结果：

numUniqCategories
shipID  
1   2
2   1
3   1
4   1
5   3

Xarray专属实现方案 #

方法一：分组后直接统计唯一值数量 #

利用groupby按shipID分组，对category变量调用nunique()方法指定维度即可完成统计：

# Xarray实现代码
result = shipDS.groupby("shipID")["category"].nunique(dim="lineItem").to_dataset(name="numUniqCategories")

输出结果：

<xarray.Dataset>
Dimensions:           (shipID: 5)
Coordinates:
  * shipID            (shipID) int64 1 2 3 4 5
Data variables:
    numUniqCategories  (shipID) int64 2 1 1 1 3

如果需要转换为Pandas DataFrame格式，执行：

result.to_dataframe()

方法二：先去重再分组计数 #

先通过分组保留每个shipID与category的唯一组合，再按shipID统计组合数量：

# 保留shipID和category的唯一组合
unique_cat_per_ship = shipDS.groupby(["shipID", "category"]).first()
# 按shipID分组统计品类数量
result = unique_cat_per_ship.groupby("shipID").count(dim="category").rename({"prodID": "numUniqCategories"})

此方法同样能得到与Pandas一致的统计结果。

内容的提问来源于stack exchange，提问作者Adam Fleischhacker