在 Polar 中流式传输时推断数据类型

Question

我最近收到了一个大于内存的镶木地板文件，其中每一列都是字符串数据类型。我想进行类型推断，将可以转换为整数和浮点数的列转换，并将正确的字符串列保留为字符串，然后沉入具有正确类型的新镶木地板文件中。关于如何做到这一点有什么建议吗？

我尝试使用内置的转换函数进行一些转换并设置

strict=False

，但这最终会将所有字符串变为空。

Answer 1

没有内置方法可以直接执行此操作。问题在于，列必须是单个数据类型，并且当您接收到文件时，极性必须在实际开始流式传输（也称为甚至查看第一行）输入数据之前知道输出模式。这是一个矛盾，因为在您读取所有数据之前，您不知道字符串到 int 或字符串到 float 的转换是否有效。

这意味着您必须对数据进行 2 次遍历。一次传递查看哪些列可以转换为哪种数据类型，然后第二次传递实际接收数据。

这是一个帮助完成第一部分的函数

def try_dtype(lf, dtype, cols_to_skip=None):
    from polars.exceptions import ComputeError
    if cols_to_skip is None:
        cols_to_check = [x for x,y in lf.schema.items() if y==pl.String]
    else:
        cols_to_check = [
            x for x,y in lf.schema.items() 
            if y==pl.String and 
            x not in cols_to_skip]
    if len(cols_to_check)==0:
        return []
    good_cols=[]
    for col in cols_to_check:
        try:
            (
                lf
                .select(pl.col(col).cast(dtype))
                .collect(streaming=True)
            )
            good_cols.append(col)
        except ComputeError:
            pass
    return good_cols

使用此功能，您可以执行以下操作：

### try int first otherwise everything will be float
int_cols = try_dtype(lf, pl.Int64)

float_cols = try_dtype(lf, pl.Float64, int_cols)

exprs = []
for col in lf.columns:
    if col in int_cols:
        exprs.append(pl.col(col).cast(pl.Int64))
    elif col in float_cols:
        exprs.append(pl.col(col).cast(pl.Float64))
    else:
        exprs.append(pl.col(col))
lf.select(*exprs).sink_parquet('final_cast.parquet')