从 Pandas 到 Polars 三十八：Polars 的“瘦身”功能

Polars 有一个内置工具来进行 dtype 瘦身。调用 shrink_dtype 表达式，它会根据列中的数据将列转换为需要最少内存量的 dtype。

shrink_dtype 是 Polars 中一个非常有用的函数，它用于优化 DataFrame 中列的数据类型，以减小内存占用并可能提高某些操作的性能。当你有一个 DataFrame，其某些列的数据类型比实际需要的更大（例如，如果一列中所有的整数都可以用一个更小的整数类型来表示）时，shrink_dtype 就可以派上用场。

在 Polars 中，你可以直接在 DataFrame 上调用 shrink_dtype() 方法来减小列的数据类型。这个方法会遍历 DataFrame 中的所有列，并尝试将它们的数据类型缩小到能够安全存储当前数据的最小类型。

import polars as pl

(
  pl.DataFrame(
   {
     "a": [1, 2, 3],
        "b": [1, 2, 2 << 32],
        "c": [-1, 2, 1 << 30],
        "d": [-112, 2, 112],
        "e": [-112, 2, 129],
        "f": ["a", "b", "c"],
        "g": [0.1, 1.32, 0.12],
        "h": [True, None, False],
     }
    )
  .select(
    pl.all().shrink_dtype()
  )
)

shape: (3, 8)
┌─────┬────────────┬────────────┬──────┬──────┬─────┬──────┬───────┐
│ a   ┆ b          ┆ c          ┆ d    ┆ e    ┆ f   ┆ g    ┆ h     │
│ --- ┆ ---        ┆ ---        ┆ ---  ┆ ---  ┆ --- ┆ ---  ┆ ---   │
│ i8  ┆ i64        ┆ i32        ┆ i8   ┆ i16  ┆ str ┆ f32  ┆ bool  │
╞═════╪════════════╪════════════╪══════╪══════╪═════╪══════╪═══════╡
│ 1   ┆ 1          ┆ -1         ┆ -112 ┆ -112 ┆ a   ┆ 0.1  ┆ true  │
│ 2   ┆ 2          ┆ 2          ┆ 2    ┆ 2    ┆ b   ┆ 1.32 ┆ null  │
│ 3   ┆ 8589934592 ┆ 1073741824 ┆ 112  ┆ 129  ┆ c   ┆ 0.12 ┆ false │
└─────┴────────────┴────────────┴──────┴──────┴─────┴──────┴───────┘

浮点数和整数默认都是64位精度。在上面来自API文档的示例中，Polars发现列“a”可以是8位，列“b”必须是64位，但列“c”可以是32位。

将数值列从64位转换为32位通常是数据科学中最简单的优化手段。内存使用量减半，计算时间也可能减少到64位的一半。

你需要检查精度的损失是否可以接受。我使用的传感器精度为0.01，所以10^-6的变化是可以接受的。

如果你转换为8位或16位，内存使用量会继续按比例下降……

……但计算时间可能不会比32位更好！

大多数现代CPU没有原生支持8位或16位，所以它们必须模拟这些类型。

具有大量重复条目的字符串列也可以有效地转换为分类类型。但这将是另一个故事了。

使用shrink_dtype时需要注意以下几点：

1. shrink_dtype() 不会改变 DataFrame 中的数据，只会改变数据的存储方式。
2. 在某些情况下，缩小数据类型可能不会带来显著的内存节省或性能提升，这取决于你的具体数据和操作。
3. 缩小数据类型后，如果后续操作导致数据超出了新类型的范围，可能会引发错误或数据丢失。因此，在使用 shrink_dtype() 之前，请确保你了解数据的特点和后续的操作需求。

另外，polars会根据你的数据的实际情况调整shrink_dtype的“瘦身”，比如：

import polars as pl

# 创建一个示例 DataFrame
df = pl.DataFrame({
    "small_ints": [0, 1, 127, 255]  # 注意：这里故意包含了 255，稍后会看到效果
})

# 查看原始数据类型
print(df.dtypes)

# 尝试缩小数据类型
df_shrunk = df.select(
    pl.all().shrink_dtype())

# 查看缩小后的数据类型
print(df_shrunk.dtypes)


[Int64]
[Int16]

在这个例子中，你可能会注意到 small_ints 列的数据类型并没有从 Int64 缩小到 Int8，而是缩小到Int16。这是因为当 shrink_dtype() 检查列中的最大值时，它发现有一个 255，这个值超出了 Int8 的范围（Int8 的范围是 -128 到 127）。因此，shrink_dtype() 可能会选择保持列的原始类型不变，或者选择一个能够安全存储当前数据的更小的类型。

  往期热门文章：

从 Pandas 到 Polars 二十六：在Polars中，不要遍历列

从 Pandas 到 Polars 二十三：如果你的数据已经排序，Polars可以为你提供助力

从 Pandas 到 Polars 十八：数据科学 2025，对未来几年内数据科学领域发展的预测或展望

从 Pandas 到 Polars 十三：流式处理的关键参数

从 Pandas 到 Polars 十：“Polars 表达式“是什么？

从 Pandas 到 Polars 六：在 Polars 中流式处理大型数据集

从 Pandas 到 Polars 0：理解Polars嵌套列类型