如何将textwrap.wrap作为ufunc应用于xarray.DataArray

你的错误源于xr.apply_ufunc的维度配置与函数输出不匹配：指定了复用band维度作为输出核心维度，但实际函数是要新增一个维度存放拆分后的子串，导致维度映射混乱触发KeyError。以下是修正后的实现：

修正后的代码 #

import xarray as xr
import numpy as np
import textwrap

def decompressor(s, chunk_len):
    # 直接处理单个字符串，返回拆分后的数组
    return np.array(textwrap.wrap(s, chunk_len))

def ufunc_decompressor(s, chunk_len):
    # 计算每个字符串拆分后的子串数量
    num_chunks = s.str.len().max().item() // chunk_len
    return xr.apply_ufunc(
        decompressor,
        s,
        chunk_len,
        output_dtypes=[np.dtype(f"U{chunk_len}")],
        input_core_dims=[[], []],  # 函数作用于单个标量元素，无核心维度
        output_core_dims=[["split"]],  # 新增split维度存放拆分结果
        dask="parallelized",
        vectorize=True,
        output_sizes={"split": num_chunks}  # 明确输出维度的长度
    )

# 调用示例
xds_split = ufunc_decompressor(xds, 3).load()

关键修改说明 #

1. 调整input_core_dims：设为[[], []]，表示函数接收的是单个字符串标量而非带维度的数组，vectorize=True会自动遍历所有元素。
2. 新增输出维度：用output_core_dims=[["split"]]声明新增一个名为split的维度，专门存放每个字符串拆分后的子串，避免复用原维度导致的冲突。
3. 明确输出维度长度：通过output_sizes指定split维度的长度，确保dask能正确计算输出数组的形状。
4. 简化decompressor函数：移除s.item()，因为vectorize=True会直接传递单个字符串给函数，无需额外提取元素。

输出效果 #

修正后，你的DataArray会从(band:1, y:2, x:2)变为(band:1, y:2, x:2, split:6)，每个位置的元素会拆分为6个长度为3的子串（如"000000000000000000"变为["000", "000", "000", "000", "000", "000"]）。

内容的提问来源于stack exchange，提问作者martin-git