Polars使用指南（一）

pandas是Python数据处理中非常经典的一个科学计算库，表形式的数据结构、丰富的API和灵活的编程语法使得pandas成为最常用的的数据分析工具。但是pandas也有一个最致命的缺陷，就是效率问题，尤其是不支持并行计算。pandas2在性能方面有了极大的提升，但是不支持并行计算依然是pandas的遗憾之一。针对这个问题，市场上也涌现出了多种解决方案，如 pandarallel、dask、ray、Pandas API on Spark 等等，亦或者是开发者基于进程池的形式自己实现并行计算，但是这些方案多会有不支持跨平台、部署麻烦、不方便调试以及和pandas API兼容性差等问题，而Polars则提供了一个综合之下最适宜的方案。

Polars除了提供API形式的访问方式之外，还可以通过SQL语法查询，本文主要介绍pl.Series相关的API，其他内容将在后续文章中介绍。

首先看一下Polars官方的介绍：
Polars is a DataFrame interface on top of an OLAP Query Engine implemented in Rust using Apache Arrow Columnar Format as the memory model.

• Lazy | eager execution
• Multi-threaded
• SIMD
• Query optimization
• Powerful expression API
• Hybrid Streaming (larger than RAM datasets)
• Rust | Python | NodeJS | R | …

我们关注几个关键词：DataFrame、OLAP查询引擎、Rust实现、Apache Arrow、内存模型、多线程，可以发现Polars同样也在尽可能保持对pandas API的语法兼容，并且底层通过rust实现，支持多线程并行计算（可以充分利用多核）。接下来我们介绍具体的函数/API，详细资料可参考Polars API官网资料。

polars需要注意点如下：

• polars修正了pandas中含有空值的整型列会被转为浮点型的问题，如果除null外都是整型，则Series也是整型；
• polars中null和NaN是不同的，np.nan是NaN；

1. 输入/输出

from datetime import datetime
import polars as pl


df = pl.DataFrame(
    {
        "integer": [1, 2, 3],
        "date": [
            datetime(2022, 1, 1),
            datetime(2022, 1, 2),
            datetime(2022, 1, 3),
        ],
        "float": [4.0, 5.0, 6.0],
    }
)
s = pl.Series("a", [1, 2, 3])

print(type(df))
# 
print(type(df.select('float')))
# 
print(type(df['float']))
# 

# 读csv
df = pl.read_csv("docs/data/output.csv")
df.write_csv(path, separator=",")
# 按照batch方式加载csv文件，通过reader.next_batches(5)依次读取文件
reader = pl.read_csv_batched("docs/data/output.csv")
# 读json
pl.read_json("docs/data/output.json")
df.write_json(row_oriented=True)
# 读parquet
pl.read_parquet("docs/data/output.parquet")
df.write_parquet(path)
# 读数据库
pl.read_database(
    query="SELECT * FROM test_data",
    connection=user_conn,
    schema_overrides={"normalised_score": pl.UInt8},
) 
# 读avro
pl.read_avro
df.write_avro(path)

2. 查询

from datetime import datetime
import polars as pl


df = pl.DataFrame(
    {
        "a": [1, 2, 3],
        "b": [
            datetime(2022, 1, 1),
            datetime(2022, 1, 2),
            datetime(2022, 1, 3),
        ],
        "c": [4.0, 5.0, 6.0],
    }
)

# 查询所有列
df.select('*')
df.select(pl.col('*'))
# 查询指定列
df.select(pl.col('a', 'b'))
df.select(['a', 'b'])
df.select(pl.col('a'), pl.col('b'))
df[['a', 'b']]
# 排除指定列
df.select(pl.exclude("a"))
# 增加虚拟列（新增字段）
df.with_columns(pl.col("b").sum().alias("e"), (pl.col("b") + 42).alias("b+42"))

3. 过滤

df.filter(pl.col("c").is_between(datetime(2022, 12, 2), datetime(2022, 12, 8)),)
df.filter((pl.col("a") <= 3) & (pl.col("d").is_not_nan()))

4. 分组

分组后利用Polars的并行计算能力也是我们非常需要的功能。

df = pl.DataFrame(
    {
        "x": range(8),
        "y": ["A", "A", "A", "B", "B", "C", "X", "X"],
    }
)

# 分组统计
df.group_by("y", maintain_order=True).count()
# 利用agg分组聚合，这里只有x一列需要聚合，所以不会有别名冲突
df.group_by("y", maintain_order=True).agg(
    pl.col("*").count().alias("count"),
    pl.col("*").sum().alias("sum"),
)

5. Series API

详情参考Series官方API。
Series的参数如下，注意和pd.Series不同，第一个参数不是data（values），但是也可以接收ArrayLike类型参数，此时不能指定name参数。

class polars.Series(
	name: str | ArrayLike | None = None,
	values: ArrayLike | None = None,
	dtype: PolarsDataType | None = None,
	*,
	strict: bool = True,
	nan_to_null: bool = False,
	dtype_if_empty: PolarsDataType = Null,
)

s1 = pl.Series("a", [1, 2, 3])
s2 = pl.Series("a", [1, 2, 3], dtype=pl.Float32)

Series常用API如下：
尤其需要注意，应优先使用表达式操作（列操作，如select、filter、with_columns、group_by），而不是map_elements / apply，因为表达式操作操作性能更高。表达式计算可以利用Rust计算、并行计算、逻辑优化，而UDF（map_elements ）往往不行。

import polars as pl

s = pl.Series([1, -2, -3])

# 绝对值
s.abs()

# rename
s.alias("b")

# and
pl.Series([False, True]).all()  # 结果False
pl.Series([None, True]).all()	# 结果True
pl.Series([None, True]).all(ignore_nulls=False)  # 结果None
# or
pl.Series([True, False]).any()  # 结果True
pl.Series([None, False]).any()  # 结果False
pl.Series([None, False]).any(ignore_nulls=False)  # 结果None

# 追加，注意会修改a，且append会同时返回a
a = pl.Series("a", [1, 2, 3])
b = pl.Series("b", [4, 5])
a.append(b)
a.n_chunks()  # 结果为2
# extend同样可实现追加功能。append的是将其他Series的chunk添加到自身（拼接），底层仍然是多个chunk。
# extend将其他Series的数据追加自身内存，因此可能会导致重新分配内存。
# 所以extend执行过程可能会比append更久，但是extend的结果会比append的结果查询更快。
# 如果是追加之后立刻查询，则建议使用extend；如果需要添加多个Series之后再查询，则建议使用append，然后再调用a.rechunk()
# Series.rechunk(*, in_place: bool = False)
a = pl.Series("a", [1, 2, 3])
b = pl.Series("b", [4, 5])
a.extend(b)
a.n_chunks()  # 结果为1

# 0.19.0以后已经被删除，改为 map_elements，参数一致
# apply，和pandas功能一致，skip_nulls为True表示空值不进入function计算，效率会更高
Series.apply(
	function: Callable[[Any], Any],
	return_dtype: PolarsDataType | None = None,
	*,
	skip_nulls: bool = True,
) → Self
# map_elements和apply效果相同，如果可以通过表达式（列操作）实现的功能（如select、filter）,应避免使用map_elements，因为表达式操作效率更高
# return_dtype 应显示指定，尤其是返回值和输入值类型不一致的情况
# 如果function的开销很大，可考虑使用@lru_cache装饰器优化
Series.map_elements(
	function: Callable[[Any], Any],
	return_dtype: PolarsDataType | None = None,
	*,
	skip_nulls: bool = True,
) → Self

# 三角函数
arccos()、arccosh()、arcsin()、arcsinh()、arctan()、arctanh()、cos()、cosh()、cot()

# arg_max、arg_min 输出是标量
s = pl.Series("a", [3, 2, 1])
s.arg_max()  # 结果为0
s.arg_min()  # 结果为2

# 排序
Series.sort(*, descending: bool = False, in_place: bool = False)
# 标记有序，对某些操作提高计算效率，如max/min
Series.set_sorted(*, descending: bool = False)

# 排序索引，输出结果是排序后对应位置对应元素的索引值，注意不是每个元素对应的排名
Series.arg_sort(
	*,
	descending: bool = False,
	nulls_last: bool = False,
) → Series
s = pl.Series("a", [5, 3, 4, 1, 2])
s.arg_sort()  # 结果是 [3 4 1 2 0]
# 获取为True的索引结果
(s == 2).arg_true()
# 获取只出现一次的值索引
s.arg_unique()
# 按索引取值
s = pl.Series("a", [1, 2, 3, 4])
s.gather([1, 3])  # 结果是[2 4]
# 按固定步长采样，每n个值取一次
Series.gather_every(n: int, offset: int = 0)

# 返回前n个，如果n小于0，表示取排除后|n|后的所有数据
Series.head(n: int = 10) → Series[source]
Series.limit(n: int = 10) → Series[source]
# 返回k个最小的元素
Series.bottom_k(k: int | IntoExprColumn = 5) → Series
# 返回k个最大的元素
Series.top_k(k: int | IntoExprColumn = 5) → Series
# 返回后n个，如果n小于0，则返回排除前|n|个后的所有数据
Series.tail(n: int = 10) → Series

# 类型转换，strict若为True 如果无法进行强制转换(例如，由于溢出)，则抛出错误。
Series.cast(
	dtype: PolarsDataType | type[int] | type[float] | type[str] | type[bool],
	*,
	strict: bool = True,
) → Self
s = pl.Series("a", [True, False, True])
s.cast(pl.UInt32)

# 计算立方根，下面两种等价
s.cbrt()
s ** (1.0 / 3)
# 计算平方根，下面两种等价
s.sqrt()
s ** 0.5

# 向上取整
Series.ceil() → Series
# 向下取整
Series.floor() → Series[source]

# 创建空数据拷贝，默认返回一个空的同类型Series，n表示需要填充几个空值，默认0，所以默认返回空，不修改原始数据s
Series.clear(n: int = 0) → Series

# 拷贝
s.clone()

# 限制边界值，小于下边界的置为下边界，大于上边界的置为上边界
# lower_bound和upper_bound可以是表达式，也可以是标量值，可以只设置一个
Series.clip(
	lower_bound: NumericLiteral | TemporalLiteral | IntoExprColumn | None = None,
	upper_bound: NumericLiteral | TemporalLiteral | IntoExprColumn | None = None,
) → Series
s.clip(1, 10)

# zip_with，类似于np.where，mask是布尔值类型的Series，如果为True，则取self对应位置的值，如果为False，则取other对应位置的值
Series.zip_with(mask: Series, other: Series) → Self
# when then otherwise，可以有多个when then，如果没写otherwise且所有条件都不满足，则返回空
df = pl.DataFrame({"foo": [1, 3, 4], "bar": [3, 4, 0]})
df.with_columns(
    pl.when(pl.col("foo") > 2)
    .then(1)
    .when(pl.col("bar") > 2)
    .then(4)
    .otherwise(-1)
    .alias("val")
)
┌─────┬─────┬─────┐
│ foo ┆ bar ┆ val │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i32 │
╞═════╪═════╪═════╡
│ 1   ┆ 3   ┆ 4   │
│ 3   ┆ 4   ┆ 1   │
│ 4   ┆ 0   ┆ 1   │
└─────┴─────┴─────┘
# 设置多个and条件
df.with_columns(
    val=pl.when(
        pl.col("bar") > 0,
        pl.col("foo") % 2 != 0,
    )
    .then(99)
    .otherwise(-1)
)
df.with_columns(val=pl.when(foo=4, bar=0).then(99).otherwise(-1))

# 统计非空元素数量
s.count()

# 计算依次累计最大值
Series.cum_max(*, reverse: bool = False)
s = pl.Series("s", [3, 5, 1])
s.cum_max()  # 结果：3 5 5
# 累乘
Series.cum_prod(*, reverse: bool = False) → Series
# 自定义累计运算
Series.cumulative_eval(
	expr: Expr,
	min_periods: int = 1,
	*,
	parallel: bool = False,
) → Series
s = pl.Series("values", [1, 2, 3, 4, 5])
s.cumulative_eval(pl.element().first() - pl.element().last() ** 2)
# 结果
[
    0.0
    -3.0
    -8.0
    -15.0
    -24.0
]

# 数据离散/切分，默认是左开右闭，left_closed=True，则设为左闭右开
Series.cut(
	breaks: Sequence[float],
	*,
	labels: Sequence[str] | None = None,
	left_closed: bool = False,
	include_breaks: bool = False,
) → Series | DataFrame
s = pl.Series("foo", [-2, -1, 0, 1, 2])
s.cut([-1, 1], labels=["a", "b", "c"])
# 结果
[
        "a"
        "a"
        "b"
        "b"
        "c"
]
# 根据分位数离散数据
Series.qcut(
	quantiles: Sequence[float] | int,
	*,
	labels: Sequence[str] | None = None,
	left_closed: bool = False,
	allow_duplicates: bool = False,
	include_breaks: bool = False,
) → Series | DataFrame

# 计算偏差，n默认为1，表示计算相邻元素之间的偏差
Series.diff(n: int = 1, null_behavior: NullBehavior = 'ignore')

# 计算内积
s1 = pl.Series("a", [1, 2, 3])
s2 = pl.Series("b", [4.0, 5.0, 6.0])
s1.dot(s2)   # 结果是32

# 删除空值，注意null和NaN不同
s = pl.Series([1.0, None, 3.0, float("nan")])
s.drop_nans()   # 结果是[1.0 null 3.0]
s.drop_nulls()  # 结果是[1.0 3.0 NaN]

# 填充空值
Series.fill_nan(value: int | float | Expr | None) → Series
# strategy{None, ‘forward’, ‘backward’, ‘min’, ‘max’, ‘mean’, ‘zero’, ‘one’}
Series.fill_null(
	value: Any | None = None,
	strategy: FillNullStrategy | None = None,
	limit: int | None = None,
) → Series
s = pl.Series("a", [1, 2, 3, None])
s.fill_null(strategy="forward")


# 指数移动加权平均，com、span、half_life、alpha之间的关系见注1
Series.ewm_mean(
	com: float | None = None,
	span: float | None = None,
	half_life: float | None = None,
	alpha: float | None = None,
	*,
	adjust: bool = True,
	min_periods: int = 1,
	ignore_nulls: bool = True,
) → Series
# 指数移动加权标准差
Series.ewm_std(
	com: float | None = None,
	span: float | None = None,
	half_life: float | None = None,
	alpha: float | None = None,
	*,
	adjust: bool = True,
	bias: bool = False,
	min_periods: int = 1,
	ignore_nulls: bool = True,
) → Series

# 指数运算
s.exp()

# sign函数
s.sign()

# 压平
s = pl.Series("a", [[1, 2, 3], [4, 5, 6]])
s.list.explode()
# 结果是 [1 2 3 4 5 6]
s = pl.Series("a", ["foo", "bar"])
s.str.explode()
# 结果是 ["f" "o" "o" "b" "a" "r"]

# 聚合，和explode相反，所有行压到一行中的一个list中
s.implode()

# 插值（空值），method {‘linear’, ‘nearest’}
Series.interpolate(method: InterpolationMethod = 'linear') → Series

# 判断是否在范围内
Series.is_between(
	lower_bound: IntoExpr,
	upper_bound: IntoExpr,
	closed: ClosedInterval = 'both',
) → Series
s.is_between(2, 4)
# 是否是重复值
s.is_duplicated() → Series
# 是否是布尔值
s.dtype == pl.Boolean
# 是否为空Series
s.is_empty()
# 是否有限值（非无穷大）
s.is_finite()
# 是否第一次出现
s.is_first_distinct()
# s1是否在s2中
s1.is_in(s2)
# 是否是NaN
s.is_nan()
# 是否是null
s.is_null()
# 是否有序
Series.is_sorted(*, descending: bool = False)

# 对数函数计算，默认以e为底
Series.log(base: float = 2.718281828459045) → Series
s.log()
# 以10为底
s.log10()
# 所有元素值+1后，做ln计算
s.log1p()

# 四舍五入
Series.round(decimals: int = 0)
# 四舍五入digits位有效数字
Series.round_sig_figs(digits: int)

# replace
Series.replace(
	old: IntoExpr | Sequence[Any] | Mapping[Any, Any],
	new: IntoExpr | Sequence[Any] | NoDefault = _NoDefault.no_default,
	*,
	default: IntoExpr | NoDefault = _NoDefault.no_default,
	return_dtype: PolarsDataType | None = None,
)
# 标量替换
s.replace(2, 100)
# 多个标量替换
s.replace([2, 3], [100, 200])
# map提换
mapping = {2: 100, 3: 200}
s.replace(mapping, default=-1)
# 若替换前后值类型不同，则最好指定return_dtype
s.replace(mapping, return_dtype=pl.UInt8)
# Series默认值
default = pl.Series([2.5, 5.0, 7.5, 10.0])
s.replace(2, 100, default=default)

# 数理统计函数(忽略空值)
s.mean()、s.median()、s.max()、s.min()、s.len()
# 如果有NaN则返回空
s.nan_max()/s.nan_min()
# 出现次数最多的值
s.mode()
# 去重，maintain_order=True表示保留原始顺序，会降低性能
Series.unique(*, maintain_order: bool = False) → Series
# 去重后元素的数量
s.n_unique()
# 每个元素出现次数，若sort=True表示按出现次数降序排序，False表示随机
Series.value_counts(*, sort: bool = False, parallel: bool = False) → DataFrame
# 分位数，interpolation：插值方法，{‘nearest’, ‘higher’, ‘lower’, ‘midpoint’, ‘linear’}
Series.quantile(
	quantile: float,
	interpolation: RollingInterpolationMethod = 'nearest',
) → float | None

# 排名
Series.rank(
	method: RankMethod = 'average',
	*,
	descending: bool = False,
	seed: int | None = None,
) → Series

# reshape
Series.reshape(dimensions: tuple[int, ...]) → Series

# 翻转
Series.reverse() → Series

# 滑动窗口，应计量避免直接使用rolling_map(效率低)，使用下面内置的rolling_xxx系列函数
Series.rolling_map(
	function: Callable[[Series], Any],
	window_size: int,
	weights: list[float] | None = None,
	min_periods: int | None = None,
	*,
	center: bool = False,
) → Series
# rolling_xxx系列函数
s.rolling_max、s.rolling_mean、s.rolling_median、s.rolling_min、s.rolling_quantile、s.rolling_skew、s.rolling_std、s.rolling_sum、s.rolling_var、

# 平移，n可以为负值，表示向上平移，fill_value 如何填充平移产生的空值
Series.shift(n: int = 1, *, fill_value: IntoExpr | None = None) → Series

# 优化内存，按实际数据适配内存，减少冗余内存（数据不再变动情况）
Series.shrink_to_fit(*, in_place: bool = False) → Series

# 计算偏度，正态分布偏度为0
Series.skew(*, bias: bool = True) → float | None

# 按索引取指定长度值，含offset对应元素
Series.slice(offset: int, length: int | None = None) → Series

# to_frame，pl.Series转为pl.DataFrame，name可以重命名字段名
Series.to_frame(name: str | None = None) → DataFrame
# to_list，use_pyarrow：使用pyarrow进行转换。
Series.to_list(*, use_pyarrow: bool | None = None)
# to_numpy，转为np.ndarray，关于to_numpy的注意事项和参数解释见注2
Series.to_numpy(
	*args: Any,
	zero_copy_only: bool = False,
	writable: bool = False,
	use_pyarrow: bool = True,
) → ndarray[Any, Any]
# to_pandas，转换为pandas.Series
Series.to_pandas(
	*args: Any,
	use_pyarrow_extension_array: bool = False,
	**kwargs: Any,
) → pd.Series[Any]
#

注1：
$$alpha=\frac{1}{1+com}\forall com\ge 0$$
$$alpha=\frac{2}{span+1}\forall span\ge 1$$
$$alpha=1-\exp \left\{\frac{-\ln (2)}{half\_life}\right\}\forall half\_life>0$$

注2：
to_numpy和to_list不同，如果Series是纯数字并且没有null（注意不是nan），则是零拷贝生成，即返回的ndarray是只读的，如果需要修改ndarray，则需要设置writable=True，表示创建一个拷贝。zero_copy_only参数表示使用零拷贝生成ndarray，但是如果需要做拷贝则会触发异常。