在Rust/Polars中对ChunkedArray类型进行抽象处理

问题 #

我尝试编写一个Trait，实现对Polars中列的每个元素/单元格应用函数并将结果添加至现有延迟数据框（lazy dataframe）的逻辑抽象。目前卡在将泛型ChunkedArray转换为Series的环节，不清楚如何满足ChunkedArray对应IntoSeries的 trait 约束。

当前代码如下：

trait MapCell: Send + Sync + Sized + 'static {
    type CellType : PolarsDataType + 'static ;
    type NativeType;
    type Builder: ChunkedBuilder<Self::NativeType, Self::CellType>;

    fn map(self, column:Expr, result_name:&str) -> anyhow::Result<Expr> {
        let output_type = GetOutput::from_type(self.datatype());

        Ok(column.map(move |col| {
            let mut builder = self.builder();
            for content in col.str()?.into_no_null_iter() {
                builder.append_value(self.apply(content));
            }
            let ca = builder.finish();
            Ok(Some(builder.finish().into_series()))
        }, output_type)
        .alias(result_name))
    }

    fn apply(&self, content: &str) -> Self::NativeType;

    fn builder(&self) -> Self::Builder;

    fn datatype(&self) -> DataType;
}

运行后出现错误：

error[E0599]: the method `into_series` exists for struct `ChunkedArray<<Self as MapCell>::CellType>`, but its trait bounds were not satisfied
   --> src/lib.rs:37:38
    |
37  |             Ok(Some(builder.finish().into_series()))
    |                                      ^^^^^^^^^^^ method cannot be called on `ChunkedArray<<Self as MapCell>::CellType>` due to unsatisfied trait bounds
    |
   ::: /Users/johsulli/.cargo/registry/src/index.crates.io-6f17d22bba15001f/polars-core-0.37.0/src/chunked_array/mod.rs:136:1
    |
136 | pub struct ChunkedArray<T: PolarsDataType> {
    | ------------------------------------------ doesn't satisfy `_: IntoSeries`
    |
   ::: /Users/johsulli/.cargo/registry/src/index.crates.io-6f17d22bba15001f/polars-core-0.37.0/src/series/implementations/mod.rs:51:1
    |
51  | pub(crate) struct SeriesWrap<T>(pub T);
    | ------------------------------- doesn't satisfy `_: SeriesTrait`
    |
    = note: the following trait bounds were not satisfied:
            `polars::series::implementations::SeriesWrap<polars::prelude::ChunkedArray<<Self as MapCell>::CellType>>: polars::prelude::SeriesTrait`
            which is required by `polars::prelude::ChunkedArray<<Self as MapCell>::CellType>: polars::prelude::IntoSeries`

由于无法访问SeriesWrap，不清楚如何满足这些约束条件，请问在此场景下该如何对ChunkedArray的类型进行抽象处理？

解决方案 #

问题核心是CellType的约束不够——ChunkedArray<T>要实现IntoSeries，除了T满足PolarsDataType，还需要满足内部隐含的SeriesTrait约束（不用管私有类型SeriesWrap，通过公开约束间接满足即可）。

具体修改步骤： #

1. 补充Trait约束
   在MapCell trait上添加where子句，直接要求ChunkedArray<Self::CellType>实现IntoSeries：

   trait MapCell: Send + Sync + Sized + 'static
   where
       ChunkedArray<Self::CellType>: IntoSeries,
   {
       type CellType : PolarsDataType + 'static ;
       type NativeType;
       type Builder: ChunkedBuilder<Self::NativeType, Self::CellType>;
       // 剩余方法不变
   }
   

2. 修复重复finish调用
   原代码里builder.finish()被调用了两次，这会直接panic（builder只能完成一次构建），改成只调用一次：

   let ca = builder.finish();
   Ok(Some(ca.into_series()))
   

3. 针对特定类型细化约束（可选）
   如果你的CellType是数值类型，可以直接用PolarsNumericType约束，它会自动满足IntoSeries要求：

   type CellType: PolarsNumericType + 'static;
   

   字符串、布尔等类型同理，用Utf8Type、BooleanType这类具体的Polars类型即可，它们都隐含满足IntoSeries的约束。

完整修改后的代码示例： #

use polars::prelude::*;
use anyhow::Result;

trait MapCell: Send + Sync + Sized + 'static
where
    ChunkedArray<Self::CellType>: IntoSeries,
{
    type CellType: PolarsDataType + 'static;
    type NativeType;
    type Builder: ChunkedBuilder<Self::NativeType, Self::CellType>;

    fn map(self, column: Expr, result_name: &str) -> Result<Expr> {
        let output_type = GetOutput::from_type(self.datatype());

        Ok(column.map(move |col| {
            let mut builder = self.builder();
            for content in col.str()?.into_no_null_iter() {
                builder.append_value(self.apply(content));
            }
            let ca = builder.finish();
            Ok(Some(ca.into_series()))
        }, output_type)
        .alias(result_name))
    }

    fn apply(&self, content: &str) -> Self::NativeType;

    fn builder(&self) -> Self::Builder;

    fn datatype(&self) -> DataType;
}

// 示例实现：将字符串转为i32
struct StrToI32;

impl MapCell for StrToI32 {
    type CellType = Int32Type;
    type NativeType = i32;
    type Builder = PrimitiveChunkedBuilder<Int32Type>;

    fn apply(&self, content: &str) -> i32 {
        content.parse().unwrap_or_default()
    }

    fn builder(&self) -> Self::Builder {
        PrimitiveChunkedBuilder::new("temp", 0, DataType::Int32)
    }

    fn datatype(&self) -> DataType {
        DataType::Int32
    }
}

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