ahri_tre_tabular/
parquet_io.rs

1use std::fs::File;
2use std::path::Path;
3
4use arrow_array::RecordBatch;
5use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
6use parquet::arrow::arrow_writer::ArrowWriter;
7use parquet::basic::{Compression, ZstdLevel};
8use parquet::file::properties::WriterProperties;
9
10use crate::{ArrowTable, Result, TabularError};
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub struct ParquetWriteOptions {
14    pub compression: Compression,
15}
16
17impl Default for ParquetWriteOptions {
18    fn default() -> Self {
19        Self {
20            compression: Compression::ZSTD(ZstdLevel::default()),
21        }
22    }
23}
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq)]
26pub struct ParquetReadOptions {
27    pub batch_size: usize,
28}
29
30impl Default for ParquetReadOptions {
31    fn default() -> Self {
32        Self { batch_size: 8192 }
33    }
34}
35
36pub fn write_parquet_file(path: impl AsRef<Path>, table: &ArrowTable) -> Result<()> {
37    write_parquet_file_with_options(path, table, ParquetWriteOptions::default())
38}
39
40pub fn write_parquet_file_with_options(
41    path: impl AsRef<Path>,
42    table: &ArrowTable,
43    options: ParquetWriteOptions,
44) -> Result<()> {
45    let file = File::create(path)?;
46    let properties = WriterProperties::builder()
47        .set_compression(options.compression)
48        .build();
49    let mut writer = ArrowWriter::try_new(file, table.schema().clone(), Some(properties))?;
50    for batch in table.batches() {
51        writer.write(batch)?;
52    }
53    writer.close()?;
54    Ok(())
55}
56
57pub fn read_parquet_file(path: impl AsRef<Path>) -> Result<ArrowTable> {
58    read_parquet_file_with_options(path, ParquetReadOptions::default())
59}
60
61pub fn read_parquet_file_with_options(
62    path: impl AsRef<Path>,
63    options: ParquetReadOptions,
64) -> Result<ArrowTable> {
65    if options.batch_size == 0 {
66        return Err(TabularError::InvalidBatchSize);
67    }
68
69    let file = File::open(path)?;
70    let builder = ParquetRecordBatchReaderBuilder::try_new(file)?;
71    let schema = builder.schema().clone();
72    let reader = builder.with_batch_size(options.batch_size).build()?;
73    let batches = reader
74        .collect::<std::result::Result<Vec<_>, _>>()?
75        .into_iter()
76        .map(|batch| RecordBatch::try_new(schema.clone(), batch.columns().to_vec()))
77        .collect::<std::result::Result<Vec<_>, _>>()?;
78    ArrowTable::try_new(schema, batches)
79}
80
81/// Reads only Parquet schema metadata, never materializes data batches.
82pub fn read_parquet_schema(path: &Path) -> Result<crate::SchemaRef> {
83    Ok(ParquetRecordBatchReaderBuilder::try_new(File::open(path)?)?
84        .schema()
85        .clone())
86}
ahri_tre_tabular/parquet_io.rs

ahri_tre_tabular/
parquet_io.rs
