如何在Polars GroupBy中为不同分组应用差异化表达式处理3D坐标

我有一个包含3D空间坐标（x、y、z）的Polars DataFrame，坐标可采用直角坐标系（R）、圆柱坐标系（C）或球坐标系（S）定义，且可能为相对于全局坐标系的局部坐标系（含平移/旋转变换）。该DataFrame包含坐标值、用于标识坐标系的CID、坐标系类型CID_TYPE，示例代码如下：

import polars as pl

df = pl.DataFrame({
        'ID' : [1, 2, 3, 4, 5, 6],
        'CID' : [1, 2, 3, 9, 10, 11],
        'CID_TYPE' : ['C', 'R', 'S', 'C', 'R', 'S'],
        'COORDS' :[
                    [1.0, 2.0, 3.0],
                    [4.0, 5.0, 6.0],
                    [7.0, 8.0, 9.0],
                    [10.0, 11.0, 12.0],
                    [13.0, 14.0, 15.0],
                    [16.0, 17.0, 18.0]
                    ]
        }
    )

需要完成三项操作：

1. 将所有圆柱/球坐标转换为直角坐标（如x=r(cosθ)、y=r(sinθ)等）；
2. 用3×3旋转矩阵与每个坐标相乘；
3. 用平移向量对每个坐标进行偏移。

平移向量和旋转矩阵存储在以CID为键的其他数据结构中（比如字典），由于数据量可达数千万至数亿行，速度至关重要，因此考虑使用Polars表达式。已知可先按CID_TYPE再按CID进行GroupBy：按CID_TYPE分组时，需对类型为'C'的坐标更新前两个坐标；按CID分组时，需将每个[X,Y,Z]坐标与对应旋转矩阵相乘并加偏移向量。可以通过遍历分组、提取至numpy数组处理后放回，但作为Polars新手，想知道如何用Polars表达式实现更高效的分组差异化处理？

步骤1：编写坐标系转换的向量化表达式 #

直接通过Polars条件表达式实现不同坐标系到直角坐标的转换，全程避免Python循环：

import polars as pl
import numpy as np

def convert_coords_expr():
    # 拆分坐标列表为单独分量
    r, theta, z_or_phi = pl.col("COORDS").list.get(0), pl.col("COORDS").list.get(1), pl.col("COORDS").list.get(2)
    
    # 圆柱坐标转直角坐标
    cartesian_cyl = pl.struct(
        x=r * pl.cos(theta),
        y=r * pl.sin(theta),
        z=z_or_phi
    ).list.to_struct()
    
    # 球坐标转直角坐标（r, theta, phi对应球坐标的半径、方位角、极角）
    cartesian_sph = pl.struct(
        x=r * pl.sin(z_or_phi) * pl.cos(theta),
        y=r * pl.sin(z_or_phi) * pl.sin(theta),
        z=r * pl.cos(z_or_phi)
    ).list.to_struct()
    
    # 根据坐标系类型选择转换逻辑，直角坐标直接保留原数据
    return pl.when(pl.col("CID_TYPE") == "C") \
             .then(cartesian_cyl) \
             .when(pl.col("CID_TYPE") == "S") \
             .then(cartesian_sph) \
             .otherwise(pl.col("COORDS")) \
             .alias("CARTESIAN_COORDS")

步骤2：关联旋转矩阵与平移向量 #

先将存储在字典中的转换参数转为Polars DataFrame，方便后续关联：

# 示例转换参数：key为CID，value是(3x3旋转矩阵, 3维平移向量)
cid_transforms = {
    1: (np.array([[1,0,0],[0,1,0],[0,0,1]]), np.array([0,0,0])),
    2: (np.array([[0,1,0],[-1,0,0],[0,0,1]]), np.array([1,2,3])),
    3: (np.array([[0,0,1],[0,1,0],[-1,0,0]]), np.array([4,5,6])),
    9: (np.array([[1,0,0],[0,0,-1],[0,1,0]]), np.array([7,8,9])),
    10: (np.array([[0,-1,0],[1,0,0],[0,0,1]]), np.array([10,11,12])),
    11: (np.array([[0,0,-1],[0,1,0],[1,0,0]]), np.array([13,14,15]))
}

# 转换为Polars DataFrame
transform_df = pl.DataFrame([
    {"CID": cid, "ROT_MAT": mat.tolist(), "TRANSLATION": vec.tolist()}
    for cid, (mat, vec) in cid_transforms.items()
])

步骤3：应用旋转变换与平移 #

通过join关联转换参数，再用表达式实现矩阵乘法和向量加法的向量化计算：

def apply_transform_expr():
    # 拆分直角坐标分量
    x, y, z = pl.col("CARTESIAN_COORDS").list.get(0), pl.col("CARTESIAN_COORDS").list.get(1), pl.col("CARTESIAN_COORDS").list.get(2)
    # 拆分旋转矩阵的行
    r1, r2, r3 = pl.col("ROT_MAT").list.get(0), pl.col("ROT_MAT").list.get(1), pl.col("ROT_MAT").list.get(2)
    # 拆分平移向量分量
    tx, ty, tz = pl.col("TRANSLATION").list.get(0), pl.col("TRANSLATION").list.get(1), pl.col("TRANSLATION").list.get(2)
    
    # 计算旋转后的坐标（矩阵乘法展开）
    rot_x = r1.list.get(0)*x + r1.list.get(1)*y + r1.list.get(2)*z
    rot_y = r2.list.get(0)*x + r2.list.get(1)*y + r2.list.get(2)*z
    rot_z = r3.list.get(0)*x + r3.list.get(1)*y + r3.list.get(2)*z
    
    # 加上平移向量得到全局坐标
    return pl.struct(
        global_x=rot_x + tx,
        global_y=rot_y + ty,
        global_z=rot_z + tz
    ).list.to_struct().alias("GLOBAL_COORDS")

# 执行完整流程
result_df = df.join(transform_df, on="CID", how="left") \
             .with_columns(convert_coords_expr()) \
             .with_columns(apply_transform_expr()) \
             .select("ID", "CID", "CID_TYPE", "COORDS", "CARTESIAN_COORDS", "GLOBAL_COORDS")

核心优势 #

• 全向量化处理：所有计算通过Polars表达式完成，利用Arrow引擎的底层优化，性能远高于遍历分组的numpy方式，适合亿级数据量。
• 避免GroupBy开销：通过join关联转换参数，无需频繁分组，减少大数据场景下的分组计算开销。
• 类型安全：利用Polars的结构化列表/Struct类型，保证坐标计算的类型一致性，降低出错概率。

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