Polars是一個使用rust開發(fā)的類似于Pandas的Dataframe庫菩咨，polars在很多地方的性能表現(xiàn)比pandas好不少，我目前嘗試在一些數(shù)據(jù)處理項目中使用polars去做熬丧。
最近在使用polars處理中文字符串長度的時候遇到一個小坑: str.lengths函數(shù)返回的是字節(jié)數(shù)而不是字符數(shù)导披。

問題復現(xiàn)

python代碼

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

輸出結果如下:

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

其中字符串"string"計算的長度6是正確的屈扎，而"字符串""得到的長度是9而不是3。
網(wǎng)上搜了一下撩匕，沒搜到相關問題（polars目前使用的人確實不多鹰晨，網(wǎng)上的討論比pandas少太多了），去github issues也沒搜到相關的問題。于是便決定自己排查一番模蜡，嫌啰嗦的同學可以直接跳到后面看問題結論漠趁。

由于polars是rust開發(fā)，而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中字符串長度的實現(xiàn)是否與它有關了卤妒。

查看源碼實現(xiàn)

先將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), 發(fā)現(xiàn)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嫉沽？？俏竞？绸硕，沒有實現(xiàn)代碼，不對呀魂毁，這樣不會報錯么玻佩？ 發(fā)現(xiàn)也沒有加@typing.overload裝飾器，那就可能是其他的地方對這個類做了修改席楚，自然就想到了python的裝飾器咬崔， 果然StringNameSpace類上有個一個裝飾器@expr_dispatch，見名知義烦秩，這個裝飾器做的應該就是將一些操作或者表達式轉發(fā)到其它地方垮斯。

下一步，查看expr_dispatch裝飾器源碼只祠，

def expr_dispatch(cls: type[T]) -> type[T]:
    # 先查看類cls(這里是： StringNameSpace) 中的屬性名稱"_accessor"的值, 這里得到namespace是"str"
    namespace = getattr(cls, "_accessor", None)
    # 然后根據(jù)namenode查找表達式實現(xiàn)
    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):
                    # 如果命名空間兜蠕，名稱和參數(shù)在表達式實現(xiàn)expr_lookup中，則覆蓋當前類型的方法
                    setattr(cls, name, call_expr(attr))
    return cls

這個裝飾器本質上就是修改被裝飾的類抛寝，將它的一些方法實現(xiàn)轉為表達式的實現(xiàn)熊杨，具體轉發(fā)細節(jié)比較繞曙旭，這里先不講了，字符串表達式的實現(xiàn)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()實現(xiàn)的夷狰，其中_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::<PySeries>().unwrap();
    m.add_class::<PyDataFrame>().unwrap();
    m.add_class::<PyLazyFrame>().unwrap();
    m.add_class::<PyLazyGroupBy>().unwrap();
    m.add_class::<dsl::PyExpr>().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實現(xiàn)了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的包裝結構體，這里通過將函數(shù)function應用到dsl::Expr中书劝，在函數(shù)function對Series進行處理进倍。上述代碼中通過ca.str_lengths()來計算字符串的長度, ca是&Utf8Chunked, Utf8Chunked是ChunkedArray<Utf8Type>的類型別名, ChunkedArray是polars的底層內存布局，polars中的數(shù)據(jù)的內存存儲格式是Arrow购对，ChunkedArray是對Arrow的封裝猾昆， Utf8Chunked實現(xiàn)了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是為了將函數(shù)string_lengths應用Utf8Chunked的每個chunked中（這里即Series的每個元素），那string_lengths就是最終我們找的代碼啦:

pub fn string_lengths(array: &Utf8Array<i64>) -> 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中骡苞，對于變長數(shù)據(jù)的存儲主要由數(shù)據(jù)數(shù)組和偏移數(shù)組構成（存儲結構示意如下）垂蜗，第個元素的長度為：offset[i + 1] - offset[i]，由于polars使用了utf8編碼字符串解幽， "string"每個字符都是英文字母贴见，每個字符占用一個字節(jié)，所以"string"的長度為6躲株， 而"字符串"中每個字符都是中文字符片部，正好這幾個中文字符每個都占用3個字節(jié)，所以長度為

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

問題結論

到這也基本清晰了霜定，polars對于中文字符串長度計算的問題档悠，主要跟polars的對字符串使用utf8編碼以及底層arrow存儲有關，與我猜測的可能是rust字符串api導致的沒有直接關系望浩。

從rust設計理念來看辖所，直接返回字符串的字節(jié)數(shù)貌似沒什么問題，畢竟rust字符串的len函數(shù)返回的就是字符串的字節(jié)數(shù)磨德，另外rust字符串直接返回字節(jié)數(shù)的時間復雜度是O(1)缘回，rust沒有直接提供獲取字符數(shù)量的api，當然也可以通過s.chars().count()獲得字符數(shù)量剖张，但是這里的時間復雜度就是O(n)了切诀。

但是從數(shù)據(jù)分析師的角度，個人認為絕大部分情況都是希望獲取字符串的長度而不是字節(jié)數(shù)搔弄，當然有一個臨時的計算方法：

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
]

這個實現(xiàn)實在丑陋且效率一般。

社區(qū)問題反饋

個人覺得可以提供一個新的api來返回字符串的長度丰滑，于是便去github提了這個issues顾犹，社區(qū)大佬立馬跟進并提了PR倒庵，很快呀，經過簡單討論炫刷，之前的str.lengthsapi不變擎宝，依然返回字符串占用的字節(jié)數(shù)，新增一個str.n_charsapi來返回字符串中字符的數(shù)量浑玛。目前最新版本的polars中已經包含了這個api绍申，所以求字符串長度可以直接使用了:

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

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

開源庫踩坑思路

總結上面的流程，我理解的踩坑思路大概是這樣：

1. 使用庫并發(fā)現(xiàn)問題
2. 搜索引擎或者項目issues等搜搜相關問題
3. 如果還無法解決顾彰，大膽猜測一下導致問題的原因极阅，可能的話做做簡單的驗證
4. 拉取庫的源碼，結合問題和猜想逐步分析并查看相關實現(xiàn)
5. issues中反饋問題
6. 根據(jù)issues的討論涨享，可以的就考慮提交PR解決相關問題

最后

經過這一番折騰筋搏，發(fā)現(xiàn)polars整體設計還是很不錯的（基于arrow的存儲設計、惰性求值和執(zhí)行計劃優(yōu)化等等）厕隧，后續(xù)有空可以再研究研究寫幾篇原理解析的文章奔脐。

另外對rust語言感興趣并想做一些項目實踐的話（沒錯，就是我啦）吁讨，polars值得一試髓迎，個人感覺polars對sql的和更多數(shù)據(jù)源的支持以及多語言api都是一些不錯的值得做的方向。