nc2parquet 0.2.0

High-performance NetCDF to Parquet converter with cloud storage support
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90

use crate::errors::Nc2ParquetError;
use crate::input::OutputConfig;
use crate::storage::{StorageBackend, StorageFactory};
use log::debug;
use polars::prelude::*;
use std::io::Cursor;

pub(crate) fn write_dataframe_to_parquet(
    df: &mut DataFrame,
    output_path: &str,
    output_config: Option<&OutputConfig>,
) -> Result<(), Nc2ParquetError> {
    debug!("Writing DataFrame to parquet file: {}\n", output_path);
    debug!("DataFrame shape: {:?}", df.shape());
    debug!("DataFrame schema:\n{:?}", df.schema());
    debug!("First few rows:\n{}", df.head(Some(5)));

    if let Some(parent) = std::path::Path::new(output_path).parent() {
        std::fs::create_dir_all(parent)?;
    }

    let file = std::fs::File::create(output_path)?;
    let writer = build_parquet_writer(file, output_config);

    writer.finish(df)?;
    debug!("Successfully wrote parquet file: {}", output_path);

    Ok(())
}

pub(crate) async fn write_dataframe_to_parquet_async(
    df: &mut DataFrame,
    output_path: &str,
    output_config: Option<&OutputConfig>,
) -> Result<(), Nc2ParquetError> {
    debug!("Writing DataFrame to parquet file: {}\n", output_path);
    debug!("DataFrame shape: {:?}", df.shape());
    debug!("DataFrame schema:\n{:?}", df.schema());
    debug!("First few rows:\n{}", df.head(Some(5)));

    let parquet_bytes = dataframe_to_parquet_bytes(df, output_config)?;
    let storage = StorageFactory::from_path(output_path).await?;
    storage.write(output_path, &parquet_bytes).await?;

    debug!("Successfully wrote parquet file: {}", output_path);
    Ok(())
}

fn dataframe_to_parquet_bytes(
    df: &mut DataFrame,
    output_config: Option<&OutputConfig>,
) -> Result<Vec<u8>, Nc2ParquetError> {
    let mut buffer = Vec::new();
    let cursor = Cursor::new(&mut buffer);
    let writer = build_parquet_writer(cursor, output_config);

    writer.finish(df)?;
    Ok(buffer)
}

/// Construct a [`ParquetWriter`] configured according to `output_config`.
///
/// When `output_config` is `None` the returned writer has no explicit
/// configuration applied (Polars defaults: Zstd, statistics enabled, no size
/// limits), preserving the identical behaviour that existed before this config
/// option was introduced.
fn build_parquet_writer<W: std::io::Write>(
    sink: W,
    output_config: Option<&OutputConfig>,
) -> ParquetWriter<W> {
    let mut writer = ParquetWriter::new(sink);

    if let Some(config) = output_config {
        writer = writer.with_compression(config.to_polars_compression());

        if let Some(rg) = config.row_group_size {
            writer = writer.with_row_group_size(Some(rg));
        }

        if let Some(dp) = config.data_page_size {
            writer = writer.with_data_page_size(Some(dp));
        }

        if !config.statistics {
            writer = writer.with_statistics(StatisticsOptions::empty());
        }
    }

    writer
}