使用 .when.then 时，Polars 中的新列如何获取其类型？

问题描述 投票：0回答：1

更新： 所有示例现在都会生成预期类型 

UInt8

极地新手，使用最新的0.18.2

我想使用 

pl.when.then.otherwise
在数据框中添加一个新列，但是当新列“继承”现有列的类型以及不继承时，我会感到困惑。

示例数据框：

import polars as pl
df = pl.DataFrame({"a": ["x", "y"], "b": [1, 2]}, schema={"a": pl.String, "b": pl.String})
df
Out:
┌─────┬─────┐
│ a   ┆ b   │
│ --- ┆ --- │
│ str ┆ u8  │
╞═════╪═════╡
│ x   ┆ 1   │
│ y   ┆ 2   │
└─────┴─────┘

这里新列的类型（和名称）与列

b
相同：

df.with_columns(
    pl.when(pl.col("a") == "x")
    .then(pl.col("b"))
    .otherwise(5)
)
Out:
┌─────┬─────┐
│ a   ┆ b   │
│ --- ┆ --- │
│ str ┆ u8  │
╞═════╪═════╡
│ x   ┆ 1   │
│ y   ┆ 5   │
└─────┴─────┘

如果我再添加一个 

when.then
，新列将获得不同的类型：

df.with_columns(
    pl.when(pl.col("a") == "x")
    .then(pl.col("b"))
    .when(pl.col("a") == "y")
    .then(0)
    .otherwise(5)
)
Out:
┌─────┬─────┐
│ a   ┆ b   │
│ --- ┆ --- │
│ str ┆ i32 │ # <-- Int32 instead of UInt8
╞═════╪═════╡
│ x   ┆ 1   │
│ y   ┆ 0   │
└─────┴─────┘

但通过另一个小更改，类型又回到了 

UInt8
，即使更改后的条件甚至不匹配任何行：

df.with_columns(
    pl.when(pl.col("a") == "x")
    .then(pl.col("b"))
    .when(pl.col("a") == "y")
    .then(0)
    .otherwise(pl.col("b")) # <-- changed to pl.col('b') instead of a number
)
Out:
┌─────┬─────┐
│ a   ┆ b   │
│ --- ┆ --- │
│ str ┆ u8  │
╞═════╪═════╡
│ x   ┆ 1   │
│ y   ┆ 0   │
└─────┴─────┘

有时，没有显式转换为 

pl.UInt8
的整数值会更改整个列的类型，有时则不会。为什么？如何确保新列的类型与以前相同？我每次都必须明确强制转换吗？

python python-polars
1个回答
0
投票

首先要注意的是，原生 python 不区分 Int 和 Float 的位大小，也不区分有符号和无符号整数，因此当您向表达式文字输入 0 和 5 时，它只是接收 int。为了避免过度解析，如果您给它一个小于 i64 的 int，那么它只会假设它应该是 i32。您可以使用一个表达式来告诉它找到最小的位大小

shrink_dtype
。即使该函数也不会采用您的 int 并使它们无符号。

另一件事需要注意的是，如果您将上次测试分配给新列，您会得到不同的行为：

df.with_columns(c=
    pl.when(pl.col("a") == "x")
    .then(pl.col("b"))
    .when(pl.col("a") == "y")
    .then(0)
    .otherwise(pl.col("b")) # <-- changed to pl.col('b') instead of a number
)
shape: (2, 3)
┌─────┬─────┬─────┐
│ a   ┆ b   ┆ c   │
│ --- ┆ --- ┆ --- │
│ str ┆ u8  ┆ i32 │
╞═════╪═════╪═════╡
│ x   ┆ 1   ┆ 1   │
│ y   ┆ 2   ┆ 0   │
└─────┴─────┴─────┘

因此，在这种情况下，使用相同的表达式，但这次分配给新列，您将获得默认的 i32。似乎对列的原始数据类型有一定的尊重，解析器将决定可以将不明确的 0 制成 u8，如您的测试中所示。

我认为问题是按照解决when/then 语句的顺序排列的，但我不确定，所以我将以解决方法结束。

解决方法

小规模：

根据您的喜好使用默认数据类型和短名称创建您自己的 

pl.lit
函数。 比如说

def u(x):
    return pl.lit(x, pl.UInt8())

然后你可以像这样将文字无符号8位整数包装在

u()
中：

df.with_columns(c=
    pl.when(pl.col("a") == "x")
    .then(pl.col("b"))
    .when(pl.col("a") == "y")
    .then(u(0))
    .otherwise(u(5))
)

如果愿意，您可以为每个数据类型创建类似的函数。

一般来说，为其他人会看到的任何东西起一个如此简短的名称是不好的做法，但如果您只是为了在实时会话中方便而这样做，那么它不会伤害任何人。

规模更大：

制作您自己的 

with_columns
替换，其目的是在替换时将结果转换回原始值，然后将其修补到 
pl.DataFrame
命名空间：

def cols(self, *kwarg, **kwargs):
    """a replacement for with_columns that attempts to cast results to original dtypes"""
    all_exprs=[]
    check_cols=[]
    for kw in kwarg:
        if kw.meta.output_name() in self.columns:
            all_exprs.append(kw.alias(f"{kw.meta.output_name()}____backup"))
            all_exprs.append(kw.cast(self.schema[kw.meta.output_name()], strict=False))
            check_cols.append(kw.meta.output_name())
        else:
            all_exprs.append(kw)            
    for key,value in kwargs.items():
        if key in self.columns:
            all_exprs.append(value.alias(f"{key}____backup"))
            all_exprs.append(value.cast(self.schema[key], strict=False))
            check_cols.append(key)
        else:
            all_exprs.append(value.alias(key))
    df=self.with_columns(all_exprs)
    revert_backup=[]
    good_to_drop=[]
    for col in check_cols:
        if df.filter(pl.col(col).is_null()!=pl.col(f"{col}____backup").is_null()).height == 0:
            good_to_drop.append(f"{col}____backup")
        else:
            revert_backup.append(col)
            good_to_drop.append(col)
    return df.drop(good_to_drop).rename({f"{x}____backup":x for x in revert_backup})
pl.DataFrame.cols=cols
del cols

其工作方式是检查所有表达式，看看它们的输出名称是否与任何现有列匹配。如果确实如此，那么它将执行该表达式两次，一次按原样执行，另一次则显式非严格强制转换回其父数据类型。然后，它会根据投射版本检查备份，看看它们的 

null
ness 是否始终相同。如果是，则意味着转换是无损的，并且会删除备份。如果该检查失败，则会删除投射版本并将备份重命名为原始名称。如果你有昂贵的表达式，那么执行两次将不是最佳的。我们可以在单独的步骤中进行转换，但是如果您只有廉价的表达式，那么添加串行步骤将不是最佳选择。我怀疑这种做法效率低下，这就是为什么这种行为尚未成为默认行为的原因。

尽管如此，这可以让你做到：

df.cols(pl.when(pl.col("a") == "x")
    .then(pl.col("b"))
    .when(pl.col("a") == "y")
    .then(1)
    .otherwise(5))

shape: (2, 2)
┌─────┬─────┐
│ a   ┆ b   │
│ --- ┆ --- │
│ str ┆ u8  │
╞═════╪═════╡
│ x   ┆ 1   │
│ y   ┆ 1   │
└─────┴─────┘

但是如果你输入 -1，例如：

df.cols(pl.when(pl.col("a") == "x")
    .then(pl.col("b"))
    .when(pl.col("a") == "y")
    .then(-1)
    .otherwise(5))

shape: (2, 2)
┌─────┬─────┐
│ a   ┆ b   │
│ --- ┆ --- │
│ str ┆ i32 │
╞═════╪═════╡
│ x   ┆ 1   │
│ y   ┆ -1  │
└─────┴─────┘

它将恢复为默认的数据类型。

最新问题
© www.soinside.com 2019 - 2024. All rights reserved.