rskit_dataset/record/
writer.rs

1//! Streaming writers for structured dataset records.
2
3use std::path::Path;
4
5use futures_util::StreamExt as _;
6use rskit_errors::{AppError, AppResult, ErrorCode};
7use serde_json::Value;
8use tokio::sync::mpsc;
9
10use super::model::DatasetRecord;
11use super::ops::BoxRecordStream;
12
13/// Streaming writer for structured dataset records.
14#[async_trait::async_trait]
15pub trait DatasetWriter: Send + Sync {
16    /// Write records to `path`, returning the number of records written.
17    async fn write(&self, records: BoxRecordStream, path: &Path) -> AppResult<usize>;
18}
19
20/// CSV writer for structured records.
21pub struct CsvWriter;
22
23#[async_trait::async_trait]
24impl DatasetWriter for CsvWriter {
25    async fn write(&self, mut records: BoxRecordStream, path: &Path) -> AppResult<usize> {
26        let (tx, mut rx) = mpsc::channel::<AppResult<DatasetRecord>>(8);
27        let path = path.to_path_buf();
28        let writer = tokio::task::spawn_blocking(move || -> AppResult<usize> {
29            let mut writer = csv::Writer::from_path(&path).map_err(|error| {
30                AppError::new(
31                    ErrorCode::Internal,
32                    format!("failed to create CSV dataset {}: {error}", path.display()),
33                )
34            })?;
35            let mut headers: Option<Vec<String>> = None;
36            let mut count = 0usize;
37            while let Some(record) = rx.blocking_recv() {
38                let record = record?;
39                let columns =
40                    headers.get_or_insert_with(|| record.fields().keys().cloned().collect());
41                if count == 0 {
42                    writer
43                        .write_record(columns.iter())
44                        .map_err(AppError::internal)?;
45                } else {
46                    ensure_csv_columns(columns, &record)?;
47                }
48                let row = columns
49                    .iter()
50                    .map(|column| {
51                        record
52                            .fields()
53                            .get(column)
54                            .map(value_to_cell)
55                            .unwrap_or_default()
56                    })
57                    .collect::<Vec<_>>();
58                writer.write_record(row).map_err(AppError::internal)?;
59                count += 1;
60            }
61            writer.flush().map_err(AppError::internal)?;
62            Ok(count)
63        });
64
65        while let Some(record) = records.next().await {
66            if tx.send(record).await.is_err() {
67                break;
68            }
69        }
70        drop(tx);
71        writer.await.map_err(AppError::internal)?
72    }
73}
74
75/// JSON Lines writer for structured records.
76pub struct JsonLinesWriter;
77
78#[async_trait::async_trait]
79impl DatasetWriter for JsonLinesWriter {
80    async fn write(&self, mut records: BoxRecordStream, path: &Path) -> AppResult<usize> {
81        let (tx, mut rx) = mpsc::channel::<AppResult<DatasetRecord>>(8);
82        let path = path.to_path_buf();
83        let writer = tokio::task::spawn_blocking(move || -> AppResult<usize> {
84            use std::io::Write as _;
85
86            let mut file = std::fs::File::create(&path).map_err(|error| {
87                AppError::new(
88                    ErrorCode::Internal,
89                    format!(
90                        "failed to create JSON Lines dataset {}: {error}",
91                        path.display()
92                    ),
93                )
94            })?;
95            let mut count = 0usize;
96            while let Some(record) = rx.blocking_recv() {
97                let json = record?.into_json();
98                serde_json::to_writer(&mut file, &json).map_err(AppError::internal)?;
99                file.write_all(b"\n").map_err(AppError::internal)?;
100                count += 1;
101            }
102            Ok(count)
103        });
104
105        while let Some(record) = records.next().await {
106            if tx.send(record).await.is_err() {
107                break;
108            }
109        }
110        drop(tx);
111        writer.await.map_err(AppError::internal)?
112    }
113}
114
115/// JSON array writer for small fixture-style datasets.
116pub struct JsonArrayWriter;
117
118#[async_trait::async_trait]
119impl DatasetWriter for JsonArrayWriter {
120    async fn write(&self, mut records: BoxRecordStream, path: &Path) -> AppResult<usize> {
121        let (tx, mut rx) = mpsc::channel::<AppResult<DatasetRecord>>(8);
122        let path = path.to_path_buf();
123        let writer = tokio::task::spawn_blocking(move || -> AppResult<usize> {
124            use std::io::Write as _;
125
126            let file = std::fs::File::create(&path).map_err(|error| {
127                AppError::new(
128                    ErrorCode::Internal,
129                    format!("failed to create JSON dataset {}: {error}", path.display()),
130                )
131            })?;
132            let mut writer = std::io::BufWriter::new(file);
133            writer.write_all(b"[").map_err(AppError::internal)?;
134
135            let mut count = 0usize;
136            while let Some(record) = rx.blocking_recv() {
137                if count > 0 {
138                    writer.write_all(b",").map_err(AppError::internal)?;
139                }
140                serde_json::to_writer(&mut writer, &record?.into_json())
141                    .map_err(AppError::internal)?;
142                count += 1;
143            }
144
145            writer.write_all(b"]").map_err(AppError::internal)?;
146            writer.flush().map_err(AppError::internal)?;
147            Ok(count)
148        });
149
150        while let Some(record) = records.next().await {
151            if tx.send(record).await.is_err() {
152                break;
153            }
154        }
155        drop(tx);
156        writer.await.map_err(AppError::internal)?
157    }
158}
159
160fn ensure_csv_columns(columns: &[String], record: &DatasetRecord) -> AppResult<()> {
161    if record.fields().len() == columns.len()
162        && columns
163            .iter()
164            .all(|column| record.fields().contains_key(column.as_str()))
165    {
166        return Ok(());
167    }
168    Err(AppError::new(
169        ErrorCode::InvalidInput,
170        format!(
171            "CSV record columns do not match established header {:?}; record has {:?}",
172            columns,
173            record.fields().keys().collect::<Vec<_>>()
174        ),
175    ))
176}
177
178fn value_to_cell(value: &Value) -> String {
179    match value {
180        Value::Null => String::new(),
181        Value::String(value) => value.clone(),
182        Value::Bool(value) => value.to_string(),
183        Value::Number(value) => value.to_string(),
184        Value::Array(_) | Value::Object(_) => value.to_string(),
185    }
186}
187
188#[cfg(test)]
189mod tests {
190    use futures::stream;
191    use serde_json::json;
192
193    use super::*;
194
195    #[test]
196    fn private_writer_helpers_validate_columns_and_convert_cells() {
197        assert!(
198            ensure_csv_columns(
199                &["a".to_string()],
200                &DatasetRecord::from_fields([("a", json!(1))])
201            )
202            .is_ok()
203        );
204        assert!(
205            ensure_csv_columns(
206                &["a".to_string()],
207                &DatasetRecord::from_fields([("b", json!(1))])
208            )
209            .is_err()
210        );
211        assert_eq!(value_to_cell(&Value::Null), "");
212        assert_eq!(value_to_cell(&json!("text")), "text");
213        assert_eq!(value_to_cell(&json!(true)), "true");
214        assert_eq!(value_to_cell(&json!(42)), "42");
215        assert_eq!(value_to_cell(&json!({"a": 1})), "{\"a\":1}");
216        assert_eq!(value_to_cell(&json!([1, 2])), "[1,2]");
217    }
218
219    #[tokio::test]
220    async fn writers_reject_directory_destinations() {
221        let dir = tempfile::tempdir().unwrap();
222        let records = || {
223            Box::pin(stream::iter(vec![Ok(DatasetRecord::from_fields([(
224                "name",
225                json!("a"),
226            )]))])) as BoxRecordStream
227        };
228
229        assert!(CsvWriter.write(records(), dir.path()).await.is_err());
230        assert!(JsonLinesWriter.write(records(), dir.path()).await.is_err());
231        assert!(JsonArrayWriter.write(records(), dir.path()).await.is_err());
232    }
233}
rskit_dataset/record/writer.rs

rskit_dataset/record/
writer.rs
