readstat/
rs_write_config.rs

1//! Output configuration for writing Arrow data to various formats.
2//!
3//! [`WriteConfig`] captures the output file path, format, compression settings,
4//! and overwrite behavior, decoupled from input path validation
5//! ([`crate::ReadStatPath`]).
6
7use std::path::{Path, PathBuf};
8
9#[cfg(feature = "parquet")]
10use parquet::basic::{BrotliLevel, Compression as ParquetCompressionCodec, GzipLevel, ZstdLevel};
11
12use crate::err::ReadStatError;
13
14/// Output file format for data conversion.
15///
16/// All variants are always present regardless of which writer features are
17/// enabled. Attempting to *write* a format whose feature is disabled returns a
18/// runtime [`ReadStatError`] from the writer rather than failing to compile.
19///
20/// This enum is `#[non_exhaustive]`: new format variants may be added in
21/// minor releases. Match with a wildcard arm to remain forward-compatible.
22#[non_exhaustive]
23#[derive(Debug, Clone, Copy)]
24pub enum OutFormat {
25    /// Comma-separated values.
26    Csv,
27    /// Feather (Arrow IPC) format.
28    Feather,
29    /// Newline-delimited JSON.
30    Ndjson,
31    /// Apache Parquet columnar format.
32    Parquet,
33}
34
35impl std::fmt::Display for OutFormat {
36    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
37        match self {
38            Self::Csv => f.write_str("csv"),
39            Self::Feather => f.write_str("feather"),
40            Self::Ndjson => f.write_str("ndjson"),
41            Self::Parquet => f.write_str("parquet"),
42        }
43    }
44}
45
46impl std::str::FromStr for OutFormat {
47    type Err = ReadStatError;
48
49    /// Parses a format name (case-insensitive) into an [`OutFormat`].
50    ///
51    /// Accepted values: `"csv"`, `"feather"`, `"ndjson"`, `"parquet"`.
52    ///
53    /// # Errors
54    ///
55    /// Returns [`ReadStatError::UnknownFormat`] for unrecognized format strings.
56    fn from_str(s: &str) -> Result<Self, Self::Err> {
57        match s.to_lowercase().as_str() {
58            "csv" => Ok(Self::Csv),
59            "feather" => Ok(Self::Feather),
60            "ndjson" => Ok(Self::Ndjson),
61            "parquet" => Ok(Self::Parquet),
62            _ => Err(ReadStatError::UnknownFormat(s.to_string())),
63        }
64    }
65}
66
67/// Parquet compression algorithm.
68///
69/// This enum is `#[non_exhaustive]`: new codec variants may be added in
70/// minor releases. Match with a wildcard arm to remain forward-compatible.
71#[non_exhaustive]
72#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
73pub enum ParquetCompression {
74    /// No compression.
75    Uncompressed,
76    /// Snappy compression (fast, moderate ratio).
77    Snappy,
78    /// Gzip compression (levels 0-9).
79    Gzip,
80    /// LZ4 raw compression.
81    Lz4Raw,
82    /// Brotli compression (levels 0-11).
83    Brotli,
84    /// Zstandard compression (levels 0-22).
85    Zstd,
86}
87
88impl std::fmt::Display for ParquetCompression {
89    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
90        match self {
91            Self::Uncompressed => f.write_str("uncompressed"),
92            Self::Snappy => f.write_str("snappy"),
93            Self::Gzip => f.write_str("gzip"),
94            Self::Lz4Raw => f.write_str("lz4-raw"),
95            Self::Brotli => f.write_str("brotli"),
96            Self::Zstd => f.write_str("zstd"),
97        }
98    }
99}
100
101impl std::str::FromStr for ParquetCompression {
102    type Err = ReadStatError;
103
104    /// Parses a codec name (case-insensitive) into a [`ParquetCompression`].
105    ///
106    /// Accepted values: `"uncompressed"`, `"snappy"`, `"gzip"`, `"lz4-raw"`
107    /// (or `"lz4raw"`), `"brotli"`, `"zstd"`.
108    ///
109    /// # Errors
110    ///
111    /// Returns [`ReadStatError::UnknownFormat`] for unrecognized codec names.
112    fn from_str(s: &str) -> Result<Self, Self::Err> {
113        match s.to_lowercase().as_str() {
114            "uncompressed" => Ok(Self::Uncompressed),
115            "snappy" => Ok(Self::Snappy),
116            "gzip" => Ok(Self::Gzip),
117            "lz4-raw" | "lz4raw" => Ok(Self::Lz4Raw),
118            "brotli" => Ok(Self::Brotli),
119            "zstd" => Ok(Self::Zstd),
120            _ => Err(ReadStatError::UnknownFormat(s.to_string())),
121        }
122    }
123}
124
125/// Output configuration for writing Arrow data.
126///
127/// Captures the output file path, format, compression settings, and overwrite
128/// behavior. Created separately from [`ReadStatPath`](crate::ReadStatPath),
129/// which handles only input path validation.
130///
131/// Fields are private and validated by the builder methods; read them via
132/// the accessor methods. This prevents constructing a config that bypasses path,
133/// extension, and compression-level validation.
134#[derive(Debug, Clone)]
135pub struct WriteConfig {
136    /// Optional output file path.
137    pub(crate) out_path: Option<PathBuf>,
138    /// Output format (defaults to CSV).
139    pub(crate) format: OutFormat,
140    /// Whether to overwrite an existing output file.
141    pub(crate) overwrite: bool,
142    /// Optional Parquet compression algorithm.
143    pub(crate) compression: Option<ParquetCompression>,
144    /// Optional Parquet compression level.
145    pub(crate) compression_level: Option<u32>,
146}
147
148impl WriteConfig {
149    /// Infers the format from an output path and returns a validated config.
150    ///
151    /// # Errors
152    ///
153    /// Returns [`ReadStatError`] if the path has an unknown extension or fails
154    /// normal output-path validation.
155    pub fn from_output(path: impl Into<PathBuf>) -> Result<Self, ReadStatError> {
156        let path = path.into();
157        let format = match path
158            .extension()
159            .and_then(|extension| extension.to_str())
160            .map(str::to_ascii_lowercase)
161            .as_deref()
162        {
163            Some("csv") => OutFormat::Csv,
164            Some("feather") => OutFormat::Feather,
165            Some("ndjson") => OutFormat::Ndjson,
166            Some("parquet") => OutFormat::Parquet,
167            _ => {
168                return Err(ReadStatError::InvalidOutputConfig(format!(
169                    "cannot infer output format from '{}'; supported extensions are .csv, .feather, .ndjson, and .parquet",
170                    path.display()
171                )));
172            }
173        };
174        Self::new(format).output(path)
175    }
176
177    /// Starts a validated configuration for `format`. CSV defaults to stdout;
178    /// other formats require [`output`](Self::output) before writer creation.
179    #[must_use]
180    pub const fn new(format: OutFormat) -> Self {
181        Self {
182            out_path: None,
183            format,
184            overwrite: false,
185            compression: None,
186            compression_level: None,
187        }
188    }
189
190    /// Sets and validates the output path.
191    pub fn output(mut self, path: impl Into<PathBuf>) -> Result<Self, ReadStatError> {
192        let path = Self::validate_out_path(Some(path.into()))?.expect("path was supplied");
193        self.out_path = Self::validate_out_extension(&path, self.format)?;
194        Ok(self)
195    }
196
197    /// Controls atomic publication when the writer is successfully finished.
198    /// When false, publication fails rather than replacing a destination that
199    /// appeared while the output was being written.
200    #[must_use]
201    pub const fn overwrite(mut self, overwrite: bool) -> Self {
202        self.overwrite = overwrite;
203        self
204    }
205
206    /// Sets and validates Parquet compression.
207    pub fn compression(
208        mut self,
209        codec: ParquetCompression,
210        level: Option<u32>,
211    ) -> Result<Self, ReadStatError> {
212        if !matches!(self.format, OutFormat::Parquet) {
213            return Err(ReadStatError::InvalidCompressionConfig(
214                "compression is only supported for Parquet".into(),
215            ));
216        }
217        self.compression_level = Self::validate_compression_level(codec, level)?;
218        self.compression = Some(codec);
219        Ok(self)
220    }
221
222    pub(crate) fn validate(&self) -> Result<(), ReadStatError> {
223        if self.out_path.is_none() && !matches!(self.format, OutFormat::Csv) {
224            return Err(ReadStatError::InvalidOutputConfig(
225                "only CSV may be written to stdout".into(),
226            ));
227        }
228        Ok(())
229    }
230
231    /// The validated output path, or `None` to write CSV to stdout.
232    #[must_use]
233    pub fn out_path(&self) -> Option<&Path> {
234        self.out_path.as_deref()
235    }
236
237    /// The output format.
238    #[must_use]
239    pub const fn format(&self) -> OutFormat {
240        self.format
241    }
242
243    /// Whether an existing output file may be overwritten.
244    #[must_use]
245    pub const fn is_overwrite(&self) -> bool {
246        self.overwrite
247    }
248
249    /// The configured Parquet compression codec, if any.
250    #[must_use]
251    pub const fn compression_codec(&self) -> Option<ParquetCompression> {
252        self.compression
253    }
254
255    /// The configured Parquet compression level, if any.
256    #[must_use]
257    pub const fn compression_level(&self) -> Option<u32> {
258        self.compression_level
259    }
260
261    /// Validates the output file extension matches the format.
262    fn validate_out_extension(
263        path: &Path,
264        format: OutFormat,
265    ) -> Result<Option<PathBuf>, ReadStatError> {
266        match path.extension().and_then(|e| e.to_str()) {
267            Some(e) if e.eq_ignore_ascii_case(&format.to_string()) => Ok(Some(path.to_owned())),
268            _ => Err(ReadStatError::OutputExtensionMismatch {
269                path: path.to_owned(),
270                expected: format.to_string(),
271            }),
272        }
273    }
274
275    /// Validates the output path exists and handles overwrite logic.
276    fn validate_out_path(path: Option<PathBuf>) -> Result<Option<PathBuf>, ReadStatError> {
277        match path {
278            None => Ok(None),
279            Some(p) => {
280                let abs_path = std::path::absolute(&p)
281                    .map_err(|e| ReadStatError::Other(format!("Failed to resolve path: {e}")))?;
282
283                match abs_path.parent() {
284                    None => Err(ReadStatError::OutputParentMissing(abs_path.clone())),
285                    Some(parent) => {
286                        if parent.exists() {
287                            Ok(Some(abs_path))
288                        } else {
289                            Err(ReadStatError::OutputParentMissing(parent.to_path_buf()))
290                        }
291                    }
292                }
293            }
294        }
295    }
296
297    /// Validates compression level is valid for the given compression algorithm.
298    fn validate_compression_level(
299        compression: ParquetCompression,
300        compression_level: Option<u32>,
301    ) -> Result<Option<u32>, ReadStatError> {
302        let (name, max_level): (&str, Option<u32>) = match compression {
303            ParquetCompression::Uncompressed => ("uncompressed", None),
304            ParquetCompression::Snappy => ("snappy", None),
305            ParquetCompression::Lz4Raw => ("lz4-raw", None),
306            ParquetCompression::Gzip => ("gzip", Some(9)),
307            ParquetCompression::Brotli => ("brotli", Some(11)),
308            ParquetCompression::Zstd => ("zstd", Some(22)),
309        };
310
311        match (max_level, compression_level) {
312            (None | Some(_), None) => Ok(None),
313            (None, Some(_)) => Err(ReadStatError::Other(format!(
314                "compression codec {name} does not support a level"
315            ))),
316            (Some(max), Some(c)) => {
317                if c <= max {
318                    Ok(Some(c))
319                } else {
320                    Err(ReadStatError::Other(format!(
321                        "The compression level of {c} is not a valid level for {name} compression. \
322                         Instead, please use values between 0-{max}."
323                    )))
324                }
325            }
326        }
327    }
328}
329
330/// Creates a uniquely named staging file beside the destination. Keeping the
331/// staging file in the same directory makes the eventual rename a same-filesystem
332/// operation.
333#[cfg(any(
334    feature = "csv",
335    feature = "feather",
336    feature = "ndjson",
337    feature = "parquet"
338))]
339pub(crate) fn open_output(config: &WriteConfig) -> Result<(std::fs::File, PathBuf), ReadStatError> {
340    create_staging_file(
341        config
342            .out_path
343            .as_ref()
344            .ok_or_else(|| ReadStatError::Other("stdout has no output file".into()))?,
345    )
346}
347
348#[cfg(any(
349    feature = "csv",
350    feature = "feather",
351    feature = "ndjson",
352    feature = "parquet"
353))]
354pub(crate) fn create_staging_file(path: &Path) -> Result<(std::fs::File, PathBuf), ReadStatError> {
355    use std::sync::atomic::{AtomicU64, Ordering};
356
357    static NEXT_STAGING_ID: AtomicU64 = AtomicU64::new(0);
358    let parent = path.parent().expect("validated output has a parent");
359    let name = path
360        .file_name()
361        .and_then(|n| n.to_str())
362        .unwrap_or("output");
363    for _ in 0..100 {
364        let id = NEXT_STAGING_ID.fetch_add(1, Ordering::Relaxed);
365        let staging = parent.join(format!(".{name}.readstat-{}-{id}.tmp", std::process::id()));
366        match std::fs::OpenOptions::new()
367            .write(true)
368            .create_new(true)
369            .open(&staging)
370        {
371            Ok(file) => return Ok((file, staging)),
372            Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => {}
373            Err(error) => return Err(error.into()),
374        }
375    }
376    Err(ReadStatError::Other(format!(
377        "could not create a unique staging file for {}",
378        path.display()
379    )))
380}
381
382#[cfg(any(
383    feature = "csv",
384    feature = "feather",
385    feature = "ndjson",
386    feature = "parquet"
387))]
388pub(crate) fn publish_staging(
389    staging: &Path,
390    destination: &Path,
391    overwrite: bool,
392) -> Result<(), ReadStatError> {
393    if overwrite {
394        #[cfg(windows)]
395        {
396            use std::os::windows::ffi::OsStrExt;
397            use windows_sys::Win32::Storage::FileSystem::{
398                MOVEFILE_REPLACE_EXISTING, MOVEFILE_WRITE_THROUGH, MoveFileExW,
399            };
400
401            let staging = staging
402                .as_os_str()
403                .encode_wide()
404                .chain(std::iter::once(0))
405                .collect::<Vec<_>>();
406            let destination = destination
407                .as_os_str()
408                .encode_wide()
409                .chain(std::iter::once(0))
410                .collect::<Vec<_>>();
411            // SAFETY: both buffers are NUL-terminated and remain alive for the
412            // duration of the call. Same-directory staging keeps this on one volume.
413            let moved = unsafe {
414                MoveFileExW(
415                    staging.as_ptr(),
416                    destination.as_ptr(),
417                    MOVEFILE_REPLACE_EXISTING | MOVEFILE_WRITE_THROUGH,
418                )
419            };
420            if moved == 0 {
421                Err(std::io::Error::last_os_error().into())
422            } else {
423                Ok(())
424            }
425        }
426        #[cfg(not(windows))]
427        {
428            std::fs::rename(staging, destination).map_err(Into::into)
429        }
430    } else {
431        std::fs::hard_link(staging, destination).map_err(|error| {
432            if error.kind() == std::io::ErrorKind::AlreadyExists {
433                ReadStatError::OutputFileExists(destination.to_owned())
434            } else {
435                error.into()
436            }
437        })?;
438        std::fs::remove_file(staging)?;
439        Ok(())
440    }
441}
442
443/// Resolves [`ParquetCompression`] and an optional level into a Parquet compression codec.
444///
445/// Defaults to Snappy when no compression is specified.
446#[cfg(feature = "parquet")]
447#[allow(clippy::cast_possible_wrap)]
448pub fn resolve_parquet_compression(
449    compression: Option<ParquetCompression>,
450    compression_level: Option<u32>,
451) -> Result<ParquetCompressionCodec, ReadStatError> {
452    let codec = match compression {
453        Some(ParquetCompression::Uncompressed) => ParquetCompressionCodec::UNCOMPRESSED,
454        Some(ParquetCompression::Snappy) | None => ParquetCompressionCodec::SNAPPY,
455        Some(ParquetCompression::Gzip) => {
456            if let Some(level) = compression_level {
457                let gzip_level = GzipLevel::try_new(level).map_err(|e| {
458                    ReadStatError::Other(format!("Invalid Gzip compression level: {e}"))
459                })?;
460                ParquetCompressionCodec::GZIP(gzip_level)
461            } else {
462                ParquetCompressionCodec::GZIP(GzipLevel::default())
463            }
464        }
465        Some(ParquetCompression::Lz4Raw) => ParquetCompressionCodec::LZ4_RAW,
466        Some(ParquetCompression::Brotli) => {
467            if let Some(level) = compression_level {
468                let brotli_level = BrotliLevel::try_new(level).map_err(|e| {
469                    ReadStatError::Other(format!("Invalid Brotli compression level: {e}"))
470                })?;
471                ParquetCompressionCodec::BROTLI(brotli_level)
472            } else {
473                ParquetCompressionCodec::BROTLI(BrotliLevel::default())
474            }
475        }
476        Some(ParquetCompression::Zstd) => {
477            if let Some(level) = compression_level {
478                let zstd_level = ZstdLevel::try_new(level as i32).map_err(|e| {
479                    ReadStatError::Other(format!("Invalid Zstd compression level: {e}"))
480                })?;
481                ParquetCompressionCodec::ZSTD(zstd_level)
482            } else {
483                ParquetCompressionCodec::ZSTD(ZstdLevel::default())
484            }
485        }
486    };
487    Ok(codec)
488}
489
490#[cfg(test)]
491mod tests {
492    use super::*;
493
494    #[test]
495    fn from_output_infers_format_case_insensitively() {
496        let config = WriteConfig::from_output("result.PARQUET").unwrap();
497        assert!(matches!(config.format(), OutFormat::Parquet));
498    }
499
500    #[test]
501    fn from_output_rejects_unknown_extension() {
502        assert!(WriteConfig::from_output("result.unknown").is_err());
503    }
504
505    // --- validate_out_extension ---
506
507    #[test]
508    fn valid_csv_out_extension() {
509        let path = Path::new("/some/output.csv");
510        let result = WriteConfig::validate_out_extension(path, OutFormat::Csv).unwrap();
511        assert!(result.is_some());
512    }
513
514    #[test]
515    fn valid_parquet_out_extension() {
516        let path = Path::new("/some/output.parquet");
517        let result = WriteConfig::validate_out_extension(path, OutFormat::Parquet).unwrap();
518        assert!(result.is_some());
519    }
520
521    #[test]
522    fn valid_feather_out_extension() {
523        let path = Path::new("/some/output.feather");
524        let result = WriteConfig::validate_out_extension(path, OutFormat::Feather).unwrap();
525        assert!(result.is_some());
526    }
527
528    #[test]
529    fn valid_ndjson_out_extension() {
530        let path = Path::new("/some/output.ndjson");
531        let result = WriteConfig::validate_out_extension(path, OutFormat::Ndjson).unwrap();
532        assert!(result.is_some());
533    }
534
535    #[test]
536    fn mismatched_out_extension() {
537        let path = Path::new("/some/output.csv");
538        assert!(WriteConfig::validate_out_extension(path, OutFormat::Parquet).is_err());
539    }
540
541    #[test]
542    fn no_out_extension() {
543        let path = Path::new("/some/output");
544        assert!(WriteConfig::validate_out_extension(path, OutFormat::Csv).is_err());
545    }
546
547    // --- validate_compression_level ---
548
549    #[test]
550    fn uncompressed_rejects_level() {
551        let result =
552            WriteConfig::validate_compression_level(ParquetCompression::Uncompressed, Some(5));
553        assert!(result.is_err());
554    }
555
556    #[test]
557    fn snappy_rejects_level() {
558        let result = WriteConfig::validate_compression_level(ParquetCompression::Snappy, Some(5));
559        assert!(result.is_err());
560    }
561
562    #[test]
563    fn lz4raw_rejects_level() {
564        let result = WriteConfig::validate_compression_level(ParquetCompression::Lz4Raw, Some(5));
565        assert!(result.is_err());
566    }
567
568    #[test]
569    fn gzip_valid_level() {
570        let result =
571            WriteConfig::validate_compression_level(ParquetCompression::Gzip, Some(5)).unwrap();
572        assert_eq!(result, Some(5));
573    }
574
575    #[test]
576    fn gzip_max_valid_level() {
577        let result =
578            WriteConfig::validate_compression_level(ParquetCompression::Gzip, Some(9)).unwrap();
579        assert_eq!(result, Some(9));
580    }
581
582    #[test]
583    fn gzip_invalid_level() {
584        assert!(
585            WriteConfig::validate_compression_level(ParquetCompression::Gzip, Some(10),).is_err()
586        );
587    }
588
589    #[test]
590    fn brotli_valid_level() {
591        let result =
592            WriteConfig::validate_compression_level(ParquetCompression::Brotli, Some(11)).unwrap();
593        assert_eq!(result, Some(11));
594    }
595
596    #[test]
597    fn brotli_invalid_level() {
598        assert!(
599            WriteConfig::validate_compression_level(ParquetCompression::Brotli, Some(12),).is_err()
600        );
601    }
602
603    #[test]
604    fn zstd_valid_level() {
605        let result =
606            WriteConfig::validate_compression_level(ParquetCompression::Zstd, Some(22)).unwrap();
607        assert_eq!(result, Some(22));
608    }
609
610    #[test]
611    fn zstd_invalid_level() {
612        assert!(
613            WriteConfig::validate_compression_level(ParquetCompression::Zstd, Some(23),).is_err()
614        );
615    }
616
617    #[test]
618    fn no_level_passes_through() {
619        let result =
620            WriteConfig::validate_compression_level(ParquetCompression::Gzip, None).unwrap();
621        assert_eq!(result, None);
622    }
623
624    // --- validate_out_path ---
625
626    #[test]
627    fn validate_out_path_none() {
628        assert!(WriteConfig::validate_out_path(None).unwrap().is_none());
629    }
630}
readstat/rs_write_config.rs

readstat/
rs_write_config.rs
