将行向量写入 Parquet 文件

Question

我知道如何将 Parquet 文件读入

Vec<Row>

。

extern crate parquet;
use parquet::file::reader::{FileReader, SerializedFileReader};
use std::{fs, sync::Arc};
use parquet::column::writer::ColumnWriter;
use parquet::{
    file::{
        properties::WriterProperties,
        writer::{FileWriter, SerializedFileWriter},
    },
    schema::parser::parse_message_type,
    schema::types::TypePtr
};
use parquet::record::Row;
use parquet::record::RowAccessor;
use std::fs::File;
use std::io::prelude::*;
use std::path::Path;
use std::path::PathBuf;

fn read_parquet(in_path: &Path) -> (Vec<Row>, TypePtr) {
    // Read Parquet input file. Return a vector of rows and the Schema
    let file = File::open(in_path).unwrap();
    let reader = SerializedFileReader::new(file).unwrap();
    let row_iter = reader.get_row_iter(None).unwrap();
    let num_rows = reader.metadata().file_metadata().num_rows();
    let rows: Vec<Row> = row_iter.collect();
    println!("num rows: {}", num_rows);

    let schema = reader.metadata().file_metadata().schema_descr().root_schema_ptr();
    (rows, schema)
}

现在，如何写出相同的数据？我正在使用 parquet 板条箱。

fn to_parquet(data: Vec<Row>, schema: TypePtr, out_path: &Path) {
    let props = Arc::new(WriterProperties::builder().build());
    let file = fs::File::create(&out_path).unwrap();
    let mut writer = SerializedFileWriter::new(file, schema, props).unwrap();
    // Now what?
}

Answer 1

这是一个有效的简单示例：

    use std::{fs, path::Path, sync::Arc};
    use parquet::{column::writer::ColumnWriter, data_type::ByteArray, file::{
        properties::WriterProperties,
        writer::{FileWriter, SerializedFileWriter},
    }, schema::parser::parse_message_type};

    #[test]
    fn sample_test() {
        let path = Path::new("./sample.parquet");

        let message_type = "
            message schema {
                REQUIRED INT32 b;
                REQUIRED BINARY msg (UTF8);
            }
        ";
        let schema = Arc::new(parse_message_type(message_type).unwrap());
        let props = Arc::new(WriterProperties::builder().build());
        let file = fs::File::create(&path).unwrap();

        let mut rows: i64 = 0;
        let data = vec![
            (10, "A"),
            (20, "B"),
            (30, "C"),
            (40, "D"),
        ];

        let mut writer = SerializedFileWriter::new(file, schema, props).unwrap();
        for (key, value) in data {
            let mut row_group_writer = writer.next_row_group().unwrap();
            let id_writer = row_group_writer.next_column().unwrap();
            if let Some(mut writer) = id_writer {
                match writer {
                    ColumnWriter::Int32ColumnWriter(ref mut typed) => {
                        let values = vec![key];
                        rows +=
                            typed.write_batch(&values[..], None, None).unwrap() as i64;
                    },
                    _ => {
                        unimplemented!();
                    }
                }
                row_group_writer.close_column(writer).unwrap();
            }
            let data_writer = row_group_writer.next_column().unwrap();
            if let Some(mut writer) = data_writer {
                match writer {
                    ColumnWriter::ByteArrayColumnWriter(ref mut typed) => {
                        let values = ByteArray::from(value);
                        rows += typed.write_batch(&[values], None, None).unwrap() as i64;
                    }
                    _ => {
                        unimplemented!();
                    }
                }
                row_group_writer.close_column(writer).unwrap();
            }
            writer.close_row_group(row_group_writer).unwrap();
        }
        writer.close().unwrap();

        println!("Wrote {}", rows);

        let bytes = fs::read(&path).unwrap();
        assert_eq!(&bytes[0..4], &[b'P', b'A', b'R', b'1']);
    }

关键是使用通过调用

RowGroupWriter

获得的

writer.next_row_group()

。

创建

sample.parquet

文件后，您可以通过运行检查其内容：

$ parquet-read ./sample.parquet
{b: 10, msg: "A"}
{b: 20, msg: "B"}
{b: 30, msg: "C"}
{b: 40, msg: "D"}

查看 arrow-rs README.md 有关

parquet-read

二进制文件的详细信息，但简而言之，您可以按如下方式获取它：

$ git clone https://github.com/apache/arrow-rs
$ cd arrow-rs/parquet
$ cargo install --path . --features cli

Answer 2

Polars 非常适合编写 Parquet 文件。请参阅 Rust 文档。

use polars::prelude::*;
use std::fs::File;

fn main() -> Result<(), PolarsError> {
    let data = vec![1, 2, 3, 4, 5]; // Vector of integers.
    let s = Series::new("data", data); // Create series from the vector.
    let mut df = DataFrame::new(vec![s])?; // Create a DataFrame from series.
    let file = File::create("example.parquet").expect("could not create file");
    ParquetWriter::new(file).finish(&mut df); // Write DataFrame to Parquet file.
    Ok(())
}

使用 Rust v1.75.0 和

Cargo.toml

:

[dependencies]
polars = { version = "0.36.2", features = ["parquet", "polars-io"] }