如何以通用方式从PostgreSQL构建Polars DataFrame？

我正尝试以通用方式从PostgreSQL读取数据到Polars DataFrame，此前参考过从MySQL读取数据到Polars的实现思路，希望无需为每个新查询单独处理列。计划将所有值按列优先顺序存储在一个大向量中，用枚举存储不同类型的值，编写了如下Rust代码：

let rows = client.query(&select_sql, &[]).await?;

#[derive(Clone, Debug)]
enum GenericValue {
    String(String),
    Bool(bool),
    DateTime(DateTime<Utc>),
    Int32(i32),
    UInt(u32),
}

let number_of_rows = rows.len();
let number_of_columns = rows.first().expect("expected one row").columns().len();

println!("Number of columns: {:?}", number_of_columns);
println!("Number of rows: {:?}", number_of_rows);

let mut data: Vec<Option<GenericValue>> = vec![None; number_of_columns * number_of_rows]; 

for (row_index, row) in rows.iter().enumerate() {
    for (col_index, column) in row.columns().iter().enumerate() {
        let colType: String = column.type_().to_string();
        
        let index = col_index * number_of_columns + row_index;

        if colType == "int4" { 
            data[index] = Some(GenericValue::Int32(row.get(col_index)));
        }
        else if colType == "text" {
            data[index] = Some(GenericValue::String(row.get(col_index)));
        }
        else if colType == "varchar" {
            data[index] = Some(GenericValue::String(row.get(col_index)));
        }
        else if colType == "bool" {
            data[index] = Some(GenericValue::Bool(row.get(col_index)));
        }
        else if colType == "timestamptz" {
            data[index] = Some(GenericValue::DateTime(row.get(col_index)));
        }
        else {
            panic!("{}", &colType);
        }
    }

    for c in 0 .. number_of_columns {
        let start_index = c * number_of_columns;
        let slice = data[start_index .. start_index + number_of_rows];

        let first: Option<GenericValue> = data[start_index];

        let chunked_data: Utf8Chunked = slice.iter().map(|v| v.unwrap()).collect();

        println!("{:?}", chunked_data);
    }
}

当前遇到的问题 #

• 不知道如何为GenericValue实现PolarsAsRef<str> trait
• 不清楚分块数据如何处理NULL值，以及如何为每个序列创建正确的分块数据
• 已放弃使用connectorx，原因是它依赖旧版Polars，且PostgreSQL分区表会破坏其并行性，无法获得性能收益

问题解答 #

1. 为GenericValue实现PolarsAsRef<str> trait #

PolarsAsRef<str>的作用是让类型能转换为&str供Polars的Utf8Chunked使用。你需要为枚举的每个变体实现字符串转换逻辑，示例如下：

use polars::prelude::PolarsAsRef;

impl PolarsAsRef<str> for GenericValue {
    fn polars_as_ref(&self) -> &str {
        match self {
            GenericValue::String(s) => s.as_str(),
            GenericValue::Bool(b) => if *b { "true" } else { "false" },
            GenericValue::DateTime(dt) => dt.to_rfc3339().as_str(),
            GenericValue::Int32(i) => i.to_string().as_str(),
            GenericValue::UInt(u) => u.to_string().as_str(),
        }
    }
}

注意：这种统一转字符串的方式会丢失原始类型信息，更合理的做法是按列实际类型创建对应ChunkedArray，而非全部转成Utf8类型。

2. 处理NULL值并创建正确的分块数据 #

当前代码中v.unwrap()会在遇到NULL值时panic，且索引计算存在错误，以下是修正后的实现：

步骤1：修复列优先存储的索引计算

列优先存储的正确索引公式应为：col_index * number_of_rows + row_index，确保同一列的数据连续存储。

步骤2：按列类型构建ChunkedArray

遍历每一列，根据列类型提取对应值并构建强类型的ChunkedArray，同时保留NULL值：

use polars::prelude::*;
use chrono::{Utc, NaiveDateTime};

// 修复索引计算，按列优先存储数据
for (row_index, row) in rows.iter().enumerate() {
    for (col_index, column) in row.columns().iter().enumerate() {
        let col_type = column.type_().to_string();
        let index = col_index * number_of_rows + row_index;

        data[index] = match col_type.as_str() {
            "int4" => Some(GenericValue::Int32(row.get(col_index))),
            "text" | "varchar" => Some(GenericValue::String(row.get(col_index))),
            "bool" => Some(GenericValue::Bool(row.get(col_index))),
            "timestamptz" => Some(GenericValue::DateTime(row.get(col_index))),
            _ => panic!("不支持的列类型: {}", col_type),
        };
    }
}

// 构建DataFrame的列集合
let mut columns = Vec::new();

for col_index in 0..number_of_columns {
    let column_meta = rows.first().unwrap().columns()[col_index];
    let col_name = column_meta.name();
    let col_type = column_meta.type_().to_string();

    // 提取当前列的所有值
    let start = col_index * number_of_rows;
    let end = start + number_of_rows;
    let col_values = &data[start..end];

    // 根据列类型创建对应ChunkedArray
    let chunked_array = match col_type.as_str() {
        "int4" => {
            let values: Vec<Option<i32>> = col_values.iter()
                .map(|opt_val| opt_val.as_ref().map(|val| match val {
                    GenericValue::Int32(i) => *i,
                    _ => panic!("列{}类型不匹配", col_name),
                }))
                .collect();
            Int32Chunked::new(col_name, values).into()
        }
        "text" | "varchar" => {
            let values: Vec<Option<&str>> = col_values.iter()
                .map(|opt_val| opt_val.as_ref().map(|val| match val {
                    GenericValue::String(s) => s.as_str(),
                    _ => panic!("列{}类型不匹配", col_name),
                }))
                .collect();
            Utf8Chunked::new(col_name, values).into()
        }
        "bool" => {
            let values: Vec<Option<bool>> = col_values.iter()
                .map(|opt_val| opt_val.as_ref().map(|val| match val {
                    GenericValue::Bool(b) => *b,
                    _ => panic!("列{}类型不匹配", col_name),
                }))
                .collect();
            BooleanChunked::new(col_name, values).into()
        }
        "timestamptz" => {
            let values: Vec<Option<NaiveDateTime>> = col_values.iter()
                .map(|opt_val| opt_val.as_ref().map(|val| match val {
                    GenericValue::DateTime(dt) => dt.naive_utc(),
                    _ => panic!("列{}类型不匹配", col_name),
                }))
                .collect();
            DateTimeChunked::new(col_name, values).into()
        }
        _ => panic!("不支持的列类型: {}", col_type),
    };

    columns.push(chunked_array);
}

// 最终构建DataFrame
let df = DataFrame::new(columns).unwrap();

关键说明

• 修复索引计算错误，确保同列数据连续，便于批量提取
• 针对每列实际类型创建强类型ChunkedArray，保留原始数据类型
• 用Option原生处理NULL值，Polars的ChunkedArray天然支持可选值输入
• 避免统一转字符串带来的性能损耗和类型丢失问题

内容的提问来源于stack exchange，提问作者Glenn Pierce