pineapple_core/io/
table.rs

1// Copyright (c) 2025, Tom Ouellette
2// Licensed under the BSD 3-Clause License
3
4use std::fs::File;
5use std::path::Path;
6
7use polars::prelude::*;
8
9use crate::error::PineappleError;
10
11/// Write a table to a CSV file
12///
13/// # Arguments
14///
15/// * `df` - A DataFrame
16/// * `output` - A string containing the name of the output file
17/// * `header` - A boolean indicating whether the output file should contain a header
18///
19/// # Examples
20///
21/// ```no_run
22/// use polars::prelude::*;
23/// use pineapple_core::io::write_table_csv;
24///
25/// let column = vec![Column::new("area".into(), [2.5, 3.1, 3.4])];
26/// let mut df: DataFrame = DataFrame::new(column).unwrap();
27///
28/// write_table_csv(&mut df, "output.csv", true).unwrap()
29/// ```
30pub fn write_table_csv<P: AsRef<Path>>(
31    df: &mut DataFrame,
32    path: P,
33    header: bool,
34) -> Result<(), PineappleError> {
35    let mut output: File = File::create(&path).map_err(|_| {
36        PineappleError::OtherError(format!(
37            "Failed to create CSV file: {}",
38            path.as_ref().to_str().unwrap()
39        ))
40    })?;
41
42    CsvWriter::new(&mut output)
43        .include_header(header)
44        .finish(df)
45        .map_err(|_| PineappleError::OtherError("Failed to write CSV file.".to_string()))
46}
47
48/// Write a table to a TSV file
49///
50/// # Arguments
51///
52/// * `df` - A DataFrame
53/// * `output` - A string containing the name of the output file
54/// * `header` - A boolean indicating whether the output file should contain a header
55///
56/// # Examples
57///
58/// ```no_run
59/// use polars::prelude::*;
60/// use pineapple_core::io::write_table_tsv;
61///
62/// let column = vec![Column::new("area".into(), [2.5, 3.1, 3.4])];
63/// let mut df: DataFrame = DataFrame::new(column).unwrap();
64///
65/// write_table_tsv(&mut df, "output.tsv", true).unwrap()
66/// ```
67pub fn write_table_tsv<P: AsRef<Path>>(
68    df: &mut DataFrame,
69    path: P,
70    header: bool,
71) -> Result<(), PineappleError> {
72    let mut output: File = File::create(&path).map_err(|_| {
73        PineappleError::OtherError(format!(
74            "Failed to create TSV file: {}",
75            path.as_ref().to_str().unwrap()
76        ))
77    })?;
78
79    CsvWriter::new(&mut output)
80        .include_header(header)
81        .with_separator("\t".as_bytes()[0])
82        .finish(df)
83        .map_err(|_| PineappleError::OtherError("Failed to write TSV file.".to_string()))
84}
85
86/// Write a table to a parquet file
87///
88/// # Arguments
89///
90/// * `df` - A DataFrame
91/// * `output` - A string containing the name of the output file
92///
93/// # Examples
94///
95/// ```no_run
96/// use polars::prelude::*;
97/// use pineapple_core::io::write_table_pq;
98///
99/// let column = vec![Column::new("area".into(), [2.5, 3.1, 3.4])];
100/// let mut df: DataFrame = DataFrame::new(column).unwrap();
101///
102/// write_table_pq(&mut df, "output.pq").unwrap()
103/// ```
104pub fn write_table_pq<P: AsRef<Path>>(df: &mut DataFrame, path: P) -> Result<(), PineappleError> {
105    let mut output: File = File::create(&path).map_err(|_| {
106        PineappleError::OtherError(format!(
107            "Failed to create TSV file: {}",
108            path.as_ref().to_str().unwrap()
109        ))
110    })?;
111
112    ParquetWriter::new(&mut output)
113        .finish(df)
114        .map(|_| ())
115        .map_err(|_| PineappleError::OtherError("Failed to write parquet file.".to_string()))
116}
117
118/// Write a DataFrame to disk
119///
120/// # Arguments
121///
122/// * `df` - A DataFrame
123/// * `output` - A string containing the name of the output file
124/// * `id` - An optional string specifying a single id to prepend as first column
125///
126/// # Examples
127///
128/// ```no_run
129/// use polars::prelude::*;
130/// use pineapple_core::io::write_table;
131///
132/// let column = vec![Column::new("area".into(), [2.5, 3.1, 3.4])];
133/// let mut df: DataFrame = DataFrame::new(column).unwrap();
134///
135/// write_table(&mut df, "output.csv").unwrap()
136/// ```
137pub fn write_table<P: AsRef<Path>>(df: &mut DataFrame, path: P) -> Result<(), PineappleError> {
138    let extension = path
139        .as_ref()
140        .extension()
141        .and_then(|s| s.to_str())
142        .map(|s| s.to_lowercase());
143
144    if let Some(ext) = extension {
145        match ext.as_str() {
146            "csv" => write_table_csv(df, path, true),
147            "tsv" => write_table_tsv(df, path, true),
148            "txt" => write_table_tsv(df, path, true),
149            "parquet" => write_table_pq(df, path),
150            "pq" => write_table_pq(df, path),
151            _ => Err(PineappleError::OtherError("Failed to write table.".to_string())),
152        }
153    } else {
154        Err(PineappleError::OtherError(
155            "Provided table path has an invalid extension. Must be one of: csv, tsv, txt, parquet, or pq.".to_string()
156        ))
157    }
158}
pineapple_core/io/table.rs

pineapple_core/io/
table.rs
