Skip to main content

f_ck/engine/
writer.rs

1use anyhow::Result;
2use polars::lazy::frame::LazyFrame;
3use polars::prelude::*;
4
5#[cfg(feature = "csv-support")]
6use polars::prelude::CsvWriter;
7
8pub struct DataWriter;
9
10impl DataWriter {
11    pub fn write_csv(df: LazyFrame, output_path: &str) -> Result<()> {
12        #[cfg(feature = "csv-support")]
13        {
14            let collected = df.collect()?;
15
16            let mut file = std::fs::File::create(output_path)?;
17            CsvWriter::new(&mut file)
18                .include_header(true)
19                .finish(&mut collected.clone())?;
20
21            Ok(())
22        }
23        #[cfg(not(feature = "csv-support"))]
24        {
25            Err(anyhow::anyhow!("CSV support not enabled for WASM build"))
26        }
27    }
28
29    pub fn write_with_format(df: LazyFrame, output_path: &str, format: &str) -> Result<()> {
30        match format.to_lowercase().as_str() {
31            "csv" => Self::write_csv(df, output_path),
32            "tsv" => Self::write_tsv(df, output_path),
33            "xlsx" => Self::write_xlsx(df, output_path),
34            "sqlite" => Self::write_sqlite(df, output_path),
35            _ => Err(anyhow::anyhow!("Unsupported output format: {}", format)),
36        }
37    }
38
39    fn write_tsv(df: LazyFrame, output_path: &str) -> Result<()> {
40        #[cfg(feature = "csv-support")]
41        {
42            let collected = df.collect()?;
43
44            let mut file = std::fs::File::create(output_path)?;
45            CsvWriter::new(&mut file)
46                .include_header(true)
47                .with_separator(b'\t')
48                .finish(&mut collected.clone())?;
49
50            Ok(())
51        }
52        #[cfg(not(feature = "csv-support"))]
53        {
54            Err(anyhow::anyhow!("TSV support not enabled for WASM build"))
55        }
56    }
57
58    fn write_xlsx(_df: LazyFrame, _output_path: &str) -> Result<()> {
59        // XLSX support would require additional dependencies like calamine
60        // For now, we'll defer this implementation
61        Err(anyhow::anyhow!("XLSX output not yet implemented"))
62    }
63
64    fn write_sqlite(_df: LazyFrame, _output_path: &str) -> Result<()> {
65        // SQLite output support would require additional SQL handling
66        // For now, we'll defer this implementation
67        Err(anyhow::anyhow!("SQLite output not yet implemented"))
68    }
69
70    pub fn preview_data(df: LazyFrame, limit: Option<usize>) -> Result<String> {
71        let mut collected = df.collect()?;
72
73        if let Some(n) = limit {
74            collected = collected.head(Some(n));
75        }
76
77        Ok(format!("{}", collected))
78    }
79}