从Polars字符串长度计算问题排查谈谈开源库踩坑思路

Polars是一个使用rust开发的类似于Pandas的Dataframe库，polars在很多地方的性能表现比pandas好不少，我目前尝试在一些数据处理项目中使用polars去做。
最近在使用polars处理中文字符串长度的时候遇到一个小坑: str.lengths函数返回的是字节数而不是字符数。

问题复现

python代码

import porars as pl
s = pl.Series(["string", "字符串"])
s.str.lengths()

输出结果如下:

shape: (2,)
Series: '' [u32]
[
    6
    9
]

其中字符串"string"计算的长度6是正确的，而"字符串""得到的长度是9而不是3。
网上搜了一下，没搜到相关问题（polars目前使用的人确实不多，网上的讨论比pandas少太多了），去github issues也没搜到相关的问题。于是便决定自己排查一番，嫌啰嗦的同学可以直接跳到后面看问题结论。

由于polars是rust开发，而rust中的字符串是使用utf8编码，所以想到问题可能出在rust字符串api上，写段rust代码测试一下:

#[test]
fn test_string_len() {
    let s1 = String::from("string");
    let s2 = String::from("字符串");
    println!("英文字符串长度: {}", s1.len());
    println!("中文字符串长度: {}", s2.len());
}

输出:

英文字符串长度: 6
中文字符串长度: 9

rust字符串api确实如此，那么接下来就是看看polars中字符串长度的实现是否与它有关了。

查看源码实现

先将polars的代码克隆到本地：

git clone https://github.com/pola-rs/polars.git

然后使用IDE或者编辑器打开它(我使用clion)

python接口代码在py-polars目录，再用pycharm打开这个目录(个人觉得pycharm提示跳转比较好，方便跟踪分析代码)。

我们前面的代码s.str.lengths()中，s是polars.Series， 故先找到它，凡是python项目，先看包的__init__.py文件，看看引用的东西都是哪里来的，这里我们先看py-polars/polars/__init__.py文件, 其中:

from polars.internals.series import Series

然后直接跳转到Series源码文件(py-polars/polars/internals/series/series.py), 发现Series是一个python的class，部分代码:

@expr_dispatch
class Series:
    @property
    def str(self) -> StringNameSpace:
        """Create an object namespace of all string related methods."""
        return StringNameSpace(self)

其中str属性方法返回的是StringNameSpace, 下一步便是查看它，StringNameSpace也是一个class, 部分代码:

@expr_dispatch
class StringNameSpace:
    """Series.str namespace."""

    _accessor = "str"

    def __init__(self, series: pli.Series):
        self._s: PySeries = series._s
    def lengths(self) -> pli.Series:

找到了其中的lengths方法，what？？？，没有实现代码，不对呀，这样不会报错么？ 发现也没有加@typing.overload装饰器，那就可能是其他的地方对这个类做了修改，自然就想到了python的装饰器， 果然StringNameSpace类上有个一个装饰器@expr_dispatch，见名知义，这个装饰器做的应该就是将一些操作或者表达式转发到其它地方。

下一步，查看expr_dispatch装饰器源码，

def expr_dispatch(cls: type[T]) -> type[T]:
    # 先查看类cls(这里是： StringNameSpace) 中的属性名称"_accessor"的值, 这里得到namespace是"str"
    namespace = getattr(cls, "_accessor", None)
    # 然后根据namenode查找表达式实现
    expr_lookup = _expr_lookup(namespace)
    for name in dir(cls):
        # 遍历类cls的方法属性等
        if not name.startswith("_"):
            attr = getattr(cls, name)
            if callable(attr):
                # 如果是一个可调用的对象（这里主要是方法）
                args = attr.__code__.co_varnames[: attr.__code__.co_argcount]
                if (namespace, name, args) in expr_lookup and _is_empty_method(attr):
                    # 如果命名空间，名称和参数在表达式实现expr_lookup中，则覆盖当前类型的方法
                    setattr(cls, name, call_expr(attr))
    return cls

这个装饰器本质上就是修改被装饰的类，将它的一些方法实现转为表达式的实现，具体转发细节比较绕，这里先不讲了，字符串表达式的实现ExprStringNameSpace在文件py-polars/polars/internals/expr/string.py中，查看代码:

class ExprStringNameSpace:
    _accessor = "str"

    def __init__(self, expr: pli.Expr):
        self._pyexpr = expr._pyexpr
 
    def lengths(self) -> pli.Expr:
        return pli.wrap_expr(self._pyexpr.str_lengths())

这里的lengths是通过调用self._pyexpr.str_lengths()实现的，其中_pyexpr对应到rust的PyExpr，polars通过pyo3在python和rust间交互, 其中py-polars模块就是一个pyo3的项目，先查看py-polars/src/lib.rs，看看polars给python暴露的模块, 部分代码：

#[pymodule]
fn polars(py: Python, m: &PyModule) -> PyResult<()> {
    ...
    m.add_class::().unwrap();
    m.add_class::().unwrap();
    m.add_class::().unwrap();
    m.add_class::().unwrap();
    m.add_class::().unwrap();
    ...
}

下一步就是跳到rust的dsl::PyExpr代码中查看(py-polars/src/lazy/dsl.rs)

#[pyclass]
#[repr(transparent)]
#[derive(Clone)]
pub struct PyExpr {
    pub inner: dsl::Expr,
}
#[pymethods]
impl PyExpr {
    pub fn str_lengths(&self) -> PyExpr {
        let function = |s: Series| {
            // 将Series转为utf8的 &Utf8Chunked
            let ca = s.utf8()?;
            // Utf8Chunked实现了Utf8NameSpaceImpl特征
            Ok(ca.str_lengths().into_series())
        };
        self.clone()
            .inner
            .map(function, GetOutput::from_type(DataType::UInt32))
            .with_fmt("str.lengths")
            .into()
    }
}

PyExpr就是dsl::Expr的包装结构体，这里通过将函数function应用到dsl::Expr中，在函数function对Series进行处理。上述代码中通过ca.str_lengths()来计算字符串的长度, ca是&Utf8Chunked, Utf8Chunked是ChunkedArray的类型别名, ChunkedArray是polars的底层内存布局，polars中的数据的内存存储格式是Arrow，ChunkedArray是对Arrow的封装， Utf8Chunked实现了Utf8NameSpaceImpl特征, Utf8NameSpaceImpl部分代码:

pub trait Utf8NameSpaceImpl: AsUtf8 {
    fn str_lengths(&self) -> UInt32Chunked {
        let ca = self.as_utf8();
        ca.apply_kernel_cast(&string_lengths)
    }
}

这里的apply_kernel_cast是为了将函数string_lengths应用Utf8Chunked的每个chunked中（这里即Series的每个元素），那string_lengths就是最终我们找的代码啦:

pub fn string_lengths(array: &Utf8Array) -> ArrayRef {
    // 通过arrow存储的偏移计算长度
    let values = array.offsets().windows(2).map(|x| (x[1] - x[0]) as u32);
    let values: Buffer<_> = Vec::from_trusted_len_iter(values).into();
    let array = UInt32Array::from_data(DataType::UInt32, values, array.validity().cloned());
    Box::new(array)
}

在arrow中，对于变长数据的存储主要由数据数组和偏移数组构成（存储结构示意如下），第个元素的长度为：offset[i + 1] - offset[i]，由于polars使用了utf8编码字符串， "string"每个字符都是英文字母，每个字符占用一个字节，所以"string"的长度为6， 而"字符串"中每个字符都是中文字符，正好这几个中文字符每个都占用3个字节，所以长度为

┌────────┬────────┐
│ data   ┆ offset │
╞════════╪════════╡
│        ┆ 0      │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ string ┆ 6      │
├╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ 字符串  ┆ 15     │
└────────┴────────┘

问题结论

到这也基本清晰了，polars对于中文字符串长度计算的问题，主要跟polars的对字符串使用utf8编码以及底层arrow存储有关，与我猜测的可能是rust字符串api导致的没有直接关系。

从rust设计理念来看，直接返回字符串的字节数貌似没什么问题，毕竟rust字符串的len函数返回的就是字符串的字节数，另外rust字符串直接返回字节数的时间复杂度是O(1)，rust没有直接提供获取字符数量的api，当然也可以通过s.chars().count()获得字符数量，但是这里的时间复杂度就是O(n)了。

但是从数据分析师的角度，个人认为绝大部分情况都是希望获取字符串的长度而不是字节数，当然有一个临时的计算方法：

import porars as pl
s = pl.Series(["string", "字符串"])
s.str.split(by="").arr.lengths().apply(lambda l: l - 2 if l >= 2 else l)

shape: (2,)
Series: '' [i64]
[
    6
    3
]

这个实现实在丑陋且效率一般。

社区问题反馈

个人觉得可以提供一个新的api来返回字符串的长度，于是便去github提了这个issues，社区大佬立马跟进并提了PR，很快呀，经过简单讨论，之前的str.lengthsapi不变，依然返回字符串占用的字节数，新增一个str.n_charsapi来返回字符串中字符的数量。目前最新版本的polars中已经包含了这个api，所以求字符串长度可以直接使用了:

import porars as pl
s = pl.Series(["string", "字符串"])
s.str.n_chars()

shape: (2,)
Series: '' [u32]
[
    6
    3
]

开源库踩坑思路

总结上面的流程，我理解的踩坑思路大概是这样：

1. 使用库并发现问题
2. 搜索引擎或者项目issues等搜搜相关问题
3. 如果还无法解决，大胆猜测一下导致问题的原因，可能的话做做简单的验证
4. 拉取库的源码，结合问题和猜想逐步分析并查看相关实现
5. issues中反馈问题
6. 根据issues的讨论，可以的就考虑提交PR解决相关问题

最后

经过这一番折腾，发现polars整体设计还是很不错的（基于arrow的存储设计、惰性求值和执行计划优化等等），后续有空可以再研究研究写几篇原理解析的文章。

另外对rust语言感兴趣并想做一些项目实践的话（没错，就是我啦），polars值得一试，个人感觉polars对sql的和更多数据源的支持以及多语言api都是一些不错的值得做的方向。