Method write

arrow-csv/src/writer.rs:241–308

Write a RecordBatch to the underlying writer

(&mut self, batch: &RecordBatch)

    /// Write a RecordBatch to the underlying writer
    pub fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
        let num_columns = batch.num_columns();
        if self.beginning {
            if self.has_headers {
                let mut headers: Vec<String> = Vec::with_capacity(num_columns);
                batch
                    .schema()
                    .fields()
                    .iter()
                    .for_each(|field| headers.push(field.name().to_string()));
                self.writer
                    .write_record(&headers[..])
                    .map_err(map_csv_error)?;
            }
            self.beginning = false;
        }

        let options = FormatOptions::default()
            .with_null(self.null_value.as_deref().unwrap_or(DEFAULT_NULL_VALUE))
            .with_date_format(self.date_format.as_deref())
            .with_datetime_format(self.datetime_format.as_deref())
            .with_timestamp_format(self.timestamp_format.as_deref())
            .with_timestamp_tz_format(self.timestamp_tz_format.as_deref())
            .with_time_format(self.time_format.as_deref());

        let converters = batch
            .columns()
            .iter()
            .map(|a| {
                if a.data_type().is_nested() {
                    Err(ArrowError::CsvError(format!(
                        "Nested type {} is not supported in CSV",
                        a.data_type()
                    )))
                } else {
                    ArrayFormatter::try_new(a.as_ref(), &options)
                }
            })
            .collect::<Result<Vec<_>, ArrowError>>()?;

        let mut buffer = String::with_capacity(1024);
        let mut byte_record = ByteRecord::with_capacity(1024, converters.len());

        for row_idx in 0..batch.num_rows() {
            byte_record.clear();
            for (col_idx, converter) in converters.iter().enumerate() {
                buffer.clear();
                converter.value(row_idx).write(&mut buffer).map_err(|e| {
                    ArrowError::CsvError(format!(
                        "Error processing row {}, col {}: {e}",
                        row_idx + 1,
                        col_idx + 1
                    ))
                })?;

                let field_bytes =
                    self.get_trimmed_field_bytes(&buffer, batch.column(col_idx).data_type());
                byte_record.push_field(field_bytes);