nordisk/graph-ar

src/io/csv.rs

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

use std::{
    fs::File,
    io::{BufWriter, Write},
    path::Path,
    sync::Arc,
};

use arrow_array::RecordBatch;
use arrow_csv::{ReaderBuilder, WriterBuilder};
use arrow_schema::Schema;

use crate::error::Result;

pub fn read(path: impl AsRef<Path>) -> Result<Vec<RecordBatch>> {
    read_typed(path, None)
}

/// Read a CSV chunk. Column types are inferred from the first 100 rows, then
/// every column named in `declared` takes the declared type and the reserved
/// `_graphAr*` index columns are int64.
pub fn read_typed(path: impl AsRef<Path>, declared: Option<&Schema>) -> Result<Vec<RecordBatch>> {
    let file = File::open(path.as_ref())?;

    let format = arrow_csv::reader::Format::default().with_header(true);
    let (inferred, _) = format.infer_schema(&mut File::open(path.as_ref())?, Some(100))?;
    let schema = super::apply_declared_types(&inferred, declared);

    let reader = ReaderBuilder::new(Arc::new(schema))
        .with_header(true)
        .build(file)?;

    let mut batches = Vec::new();
    for batch in reader {
        batches.push(batch?);
    }
    Ok(batches)
}

pub fn write(path: impl AsRef<Path>, batches: &[RecordBatch]) -> Result<()> {
    if batches.is_empty() {
        return Ok(());
    }
    let file = BufWriter::new(File::create(path)?);
    write_batches_to(file, batches)
}

/// Write `batches` as CSV into an arbitrary writer, then **flush it
/// explicitly** so a flush error (e.g. ENOSPC) surfaces here rather than being
/// silently swallowed when the buffer is dropped — mirroring `json::write`.
/// Without the explicit flush a partial/failed CSV write could be reported as
/// success. Generic over the writer so the flush-error path is unit-testable.
fn write_batches_to<W: Write>(w: W, batches: &[RecordBatch]) -> Result<()> {
    let mut writer = WriterBuilder::new().with_header(true).build(w);
    for batch in batches {
        writer.write(batch)?;
    }
    writer.into_inner().flush()?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use arrow_array::Int64Array;
    use arrow_schema::{DataType, Field, Schema};
    use std::io::{self, Write};

    /// A writer that accepts all bytes but always fails on `flush`, emulating a
    /// device that errors only when the buffer is drained (e.g. ENOSPC).
    struct FlushFails;
    impl Write for FlushFails {
        fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
            Ok(buf.len())
        }
        fn flush(&mut self) -> io::Result<()> {
            Err(io::Error::other("flush failed"))
        }
    }

    fn sample() -> Vec<RecordBatch> {
        let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)]));
        vec![RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(vec![1, 2]))]).unwrap()]
    }

    #[test]
    fn flush_error_is_propagated_not_swallowed() {
        // FAIL-ON-BUG: the old code never flushed explicitly, so the drop-time
        // flush error was discarded and this returned Ok. PASS-ON-FIX: the
        // explicit `into_inner().flush()?` propagates the error.
        let res = write_batches_to(FlushFails, &sample());
        assert!(
            res.is_err(),
            "flush failure must propagate as an error, got {res:?}"
        );
    }

    #[test]
    fn successful_write_still_ok() {
        // A writer that succeeds on flush still returns Ok (no regression).
        let res = write_batches_to(Vec::<u8>::new(), &sample());
        assert!(res.is_ok());
    }
}