如何使用xarray实现数据变量值的左连接并计算货运总重量

问题背景 #

现有两个Xarray数据集：

1. 货运数据集shipDS：每个shipID可对应多个产品，记录了产品ID和发货数量

In [2]: shipDS
Out[2]:
<xarray.Dataset>
Dimensions:     (shipID: 6)
Coordinates:
  * shipID      (shipID) int64 1 1 2 3 4 4
Data variables:
    prodID      (shipID) int64 90 91 92 90 90 91
    qtyShipped  (shipID) int64 1 1 1 1 1 1

2. 产品属性数据集prodDS：记录了每个产品的名称和重量

In [3]: prodDS
Out[3]:
<xarray.Dataset>
Dimensions:     (prodID: 4)
Coordinates:
  * prodID      (prodID) int64 90 91 92 93
Data variables:
    prodName    (prodID) <U8 'Almonds' 'Berries' 'Candy' 'Dog Food'
    prodWeight  (prodID) float64 0.5 1.5 1.0 12.0

核心需求：不切换到Pandas，仅使用Xarray计算每个货运的总重量。

数据集复现代码 #

import xarray as xr

## 构建货运数据集：每个货运可能包含多个产品
shipIDs = [1,1,2,3,4,4]
productIDs = [90,91,92,90,90,91]
quantity = [1,1,1,1,1,1]

shipDS = xr.Dataset(data_vars = 
    {
        "prodID": ("shipID", productIDs),
        "qtyShipped": ("shipID", quantity)
    },
    coords = {
        "shipID": shipIDs
    }
)

## 构建产品属性数据集
prodID = [90,91,92,93]
prodName = ["Almonds","Berries","Candy","Dog Food"]
prodWeight = [0.5,1.5,1.0,12]

prodDS = xr.Dataset(data_vars = 
    {
        "prodName": ("prodID", prodName),
        "prodWeight": ("prodID", prodWeight)
    },
    coords = {
        "prodID": prodID
    }
)

用户尝试的方法及问题 #

用户曾尝试使用merge方法：

shipDS.set_coords("prodID").merge(prodDS).to_dataframe()

但该方法需要将prodID设为坐标，导致数据空间过度膨胀，还引入了未发货的产品（如Dog Food），不符合需求。

Pandas参考解法 #

In [5]: ## Pandas实现方案
   ...: (
   ...:     shipDS   # 从货运数据集开始
   ...:     .to_dataframe()  # 转换为Pandas DataFrame
   ...:     .join(prodDS.to_dataframe(), on="prodID")   # 左连接产品属性数据集
   ...:     .reset_index()  # 重置索引
   ...:     .groupby("shipID")  # 按shipID分组
   ...:     .agg(totalWt = ("prodWeight", sum))  # 计算总重量
   ...: )
Out[5]:
        totalWt
shipID
1           2.0
2           1.0
3           0.5
4           2.0

Xarray纯实现方案 #

步骤1：匹配产品重量 #

利用Xarray的索引匹配功能，将prodDS中的产品重量关联到shipDS对应的产品上：

# 将对应产品的重量添加到货运数据集
shipDS["prodWeight"] = prodDS.prodWeight.sel(prodID=shipDS.prodID)

步骤2：计算单个产品的总重量 #

用发货数量乘以产品重量，得到每个产品在货运中的总重量：

shipDS["totalProdWeight"] = shipDS["qtyShipped"] * shipDS["prodWeight"]

步骤3：按货运ID分组求和 #

通过groupby按shipID分组，对每个分组的产品总重量求和：

total_ship_weight = shipDS.groupby("shipID").sum("shipID")["totalProdWeight"]

最终结果 #

In [x]: total_ship_weight
Out[x]:
<xarray.DataArray 'totalProdWeight' (shipID: 4)>
array([2. , 1. , 0.5, 2. ])
Coordinates:
  * shipID    (shipID) int64 1 2 3 4

结果与Pandas解法完全一致，且全程未切换到Pandas。

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