common_datasource/file_format/
csv.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::HashMap;
16use std::io;
17use std::str::FromStr;
18use std::sync::Arc;
19use std::task::Poll;
20
21use arrow::csv::reader::Format;
22use arrow::csv::{self, WriterBuilder};
23use arrow::error::ArrowError;
24use arrow::record_batch::RecordBatch;
25use arrow_schema::{Schema, SchemaRef};
26use async_trait::async_trait;
27use bytes::{Buf, Bytes};
28use common_runtime;
29use common_telemetry::warn;
30use datafusion::physical_plan::SendableRecordBatchStream;
31use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
32use futures::StreamExt;
33use futures::stream::BoxStream;
34use object_store::ObjectStore;
35use snafu::ResultExt;
36use tokio_util::compat::FuturesAsyncReadCompatExt;
37use tokio_util::io::SyncIoBridge;
38
39use crate::buffered_writer::DfRecordBatchEncoder;
40use crate::compression::CompressionType;
41use crate::error::{self, Result};
42use crate::file_format::{self, FileFormat, stream_to_file};
43use crate::share_buffer::SharedBuffer;
44use crate::util::normalize_infer_schema;
45
46const SKIP_BAD_RECORDS_BATCH_SIZE: usize = 1;
47
48#[derive(Debug, Clone, PartialEq, Eq)]
49pub struct CsvFormat {
50    pub has_header: bool,
51    pub skip_bad_records: bool,
52    pub strict_headers: bool,
53    pub delimiter: u8,
54    pub schema_infer_max_record: Option<usize>,
55    pub compression_type: CompressionType,
56    pub timestamp_format: Option<String>,
57    pub time_format: Option<String>,
58    pub date_format: Option<String>,
59}
60
61impl TryFrom<&HashMap<String, String>> for CsvFormat {
62    type Error = error::Error;
63
64    fn try_from(value: &HashMap<String, String>) -> Result<Self> {
65        let mut format = CsvFormat::default();
66        if let Some(delimiter) = value.get(file_format::FORMAT_DELIMITER) {
67            // TODO(weny): considers to support parse like "\t" (not only b'\t')
68            format.delimiter = u8::from_str(delimiter).map_err(|_| {
69                error::ParseFormatSnafu {
70                    key: file_format::FORMAT_DELIMITER,
71                    value: delimiter,
72                }
73                .build()
74            })?;
75        };
76        if let Some(compression_type) = value.get(file_format::FORMAT_COMPRESSION_TYPE) {
77            format.compression_type = CompressionType::from_str(compression_type)?;
78        };
79        if let Some(schema_infer_max_record) =
80            value.get(file_format::FORMAT_SCHEMA_INFER_MAX_RECORD)
81        {
82            format.schema_infer_max_record =
83                Some(schema_infer_max_record.parse::<usize>().map_err(|_| {
84                    error::ParseFormatSnafu {
85                        key: file_format::FORMAT_SCHEMA_INFER_MAX_RECORD,
86                        value: schema_infer_max_record,
87                    }
88                    .build()
89                })?);
90        };
91        let headers = value
92            .get(file_format::FORMAT_HEADERS)
93            .map(|headers| parse_bool(file_format::FORMAT_HEADERS, headers))
94            .transpose()?;
95        let has_header = value
96            .get(file_format::FORMAT_HAS_HEADER)
97            .map(|has_header| parse_bool(file_format::FORMAT_HAS_HEADER, has_header))
98            .transpose()?;
99        match (headers, has_header) {
100            (Some(headers), Some(has_header)) if headers != has_header => {
101                return error::ParseFormatSnafu {
102                    key: file_format::FORMAT_HEADERS,
103                    value: format!("headers={headers}, has_header={has_header}"),
104                }
105                .fail();
106            }
107            (Some(headers), _) => format.has_header = headers,
108            (_, Some(has_header)) => format.has_header = has_header,
109            _ => {}
110        }
111        if let Some(skip_bad_records) = value.get(file_format::FORMAT_SKIP_BAD_RECORDS) {
112            format.skip_bad_records =
113                parse_bool(file_format::FORMAT_SKIP_BAD_RECORDS, skip_bad_records)?;
114        };
115        if let Some(strict_headers) = value.get(file_format::FORMAT_STRICT_HEADERS) {
116            format.strict_headers = parse_bool(file_format::FORMAT_STRICT_HEADERS, strict_headers)?;
117        };
118        if format.strict_headers && !format.has_header {
119            return error::ParseFormatSnafu {
120                key: file_format::FORMAT_STRICT_HEADERS,
121                value: "strict_headers=true requires headers=true",
122            }
123            .fail();
124        }
125        if let Some(timestamp_format) = value.get(file_format::TIMESTAMP_FORMAT) {
126            format.timestamp_format = Some(timestamp_format.clone());
127        }
128        if let Some(time_format) = value.get(file_format::TIME_FORMAT) {
129            format.time_format = Some(time_format.clone());
130        }
131        if let Some(date_format) = value.get(file_format::DATE_FORMAT) {
132            format.date_format = Some(date_format.clone());
133        }
134        Ok(format)
135    }
136}
137
138fn parse_bool(key: &'static str, value: &str) -> Result<bool> {
139    value
140        .parse()
141        .map_err(|_| error::ParseFormatSnafu { key, value }.build())
142}
143
144impl Default for CsvFormat {
145    fn default() -> Self {
146        Self {
147            has_header: true,
148            skip_bad_records: false,
149            strict_headers: false,
150            delimiter: b',',
151            schema_infer_max_record: Some(file_format::DEFAULT_SCHEMA_INFER_MAX_RECORD),
152            compression_type: CompressionType::Uncompressed,
153            timestamp_format: None,
154            time_format: None,
155            date_format: None,
156        }
157    }
158}
159
160#[async_trait]
161impl FileFormat for CsvFormat {
162    async fn infer_schema(&self, store: &ObjectStore, path: &str) -> Result<Schema> {
163        let meta = store
164            .stat(path)
165            .await
166            .context(error::ReadObjectSnafu { path })?;
167
168        let reader = store
169            .reader(path)
170            .await
171            .context(error::ReadObjectSnafu { path })?
172            .into_futures_async_read(0..meta.content_length())
173            .await
174            .context(error::ReadObjectSnafu { path })?
175            .compat();
176
177        let decoded = self.compression_type.convert_async_read(reader);
178
179        let delimiter = self.delimiter;
180        let schema_infer_max_record = self.schema_infer_max_record;
181        let has_header = self.has_header;
182
183        common_runtime::spawn_blocking_global(move || {
184            let reader = SyncIoBridge::new(decoded);
185
186            let format = Format::default()
187                .with_delimiter(delimiter)
188                .with_header(has_header);
189            let (schema, _records_read) = format
190                .infer_schema(reader, schema_infer_max_record)
191                .context(error::InferSchemaSnafu)?;
192
193            Ok(normalize_infer_schema(schema))
194        })
195        .await
196        .context(error::JoinHandleSnafu)?
197    }
198}
199
200pub async fn stream_to_csv(
201    stream: SendableRecordBatchStream,
202    store: ObjectStore,
203    path: &str,
204    threshold: usize,
205    concurrency: usize,
206    format: &CsvFormat,
207) -> Result<usize> {
208    stream_to_file(
209        stream,
210        store,
211        path,
212        threshold,
213        concurrency,
214        format.compression_type,
215        |buffer| {
216            let mut builder = WriterBuilder::new();
217            if let Some(timestamp_format) = &format.timestamp_format {
218                builder = builder.with_timestamp_format(timestamp_format.to_owned())
219            }
220            if let Some(date_format) = &format.date_format {
221                builder = builder.with_date_format(date_format.to_owned())
222            }
223            if let Some(time_format) = &format.time_format {
224                builder = builder.with_time_format(time_format.to_owned())
225            }
226            builder.build(buffer)
227        },
228    )
229    .await
230}
231
232impl DfRecordBatchEncoder for csv::Writer<SharedBuffer> {
233    fn write(&mut self, batch: &RecordBatch) -> Result<()> {
234        self.write(batch).context(error::WriteRecordBatchSnafu)
235    }
236}
237
238/// Builds a CSV stream that can skip selected record-level parse/cast errors.
239///
240/// This recovery path intentionally uses one-record batches. It is slower than
241/// normal CSV scanning, but keeps each parse/cast failure isolated to a single
242/// record. Arrow's CSV decoder clears buffered rows before type parsing, so a
243/// failed multi-row flush cannot be safely retried row by row without replaying
244/// input bytes.
245pub async fn tolerant_csv_stream(
246    store: &ObjectStore,
247    path: &str,
248    schema: SchemaRef,
249    projection: Vec<usize>,
250    format: &CsvFormat,
251) -> Result<SendableRecordBatchStream> {
252    let meta = store
253        .stat(path)
254        .await
255        .context(error::ReadObjectSnafu { path })?;
256
257    let reader = store
258        .reader(path)
259        .await
260        .context(error::ReadObjectSnafu { path })?
261        .into_bytes_stream(0..meta.content_length())
262        .await
263        .context(error::ReadObjectSnafu { path })?;
264
265    let reader = format.compression_type.convert_stream(reader).boxed();
266    tolerant_csv_stream_from_reader(
267        reader,
268        path,
269        schema,
270        projection,
271        format.has_header,
272        format.delimiter,
273    )
274}
275
276fn tolerant_csv_stream_from_reader(
277    reader: BoxStream<'static, io::Result<Bytes>>,
278    path: &str,
279    schema: SchemaRef,
280    projection: Vec<usize>,
281    has_header: bool,
282    delimiter: u8,
283) -> Result<SendableRecordBatchStream> {
284    let projected_schema = Arc::new(
285        schema
286            .project(&projection)
287            .context(error::InferSchemaSnafu)?,
288    );
289    let mut decoder = csv::ReaderBuilder::new(schema)
290        .with_header(has_header)
291        .with_delimiter(delimiter)
292        .with_batch_size(SKIP_BAD_RECORDS_BATCH_SIZE)
293        .with_projection(projection)
294        .build_decoder();
295
296    let path = path.to_string();
297    let mut upstream = reader.fuse();
298    let mut buffered = Bytes::new();
299    let mut input_finished = false;
300    let stream = futures::stream::poll_fn(move |cx| {
301        loop {
302            while !input_finished {
303                if buffered.is_empty() {
304                    match futures::ready!(upstream.poll_next_unpin(cx)) {
305                        Some(Ok(bytes)) if bytes.is_empty() => continue,
306                        Some(Ok(bytes)) => buffered = bytes,
307                        Some(Err(error)) => return Poll::Ready(Some(Err(error.into()))),
308                        None => input_finished = true,
309                    }
310                }
311
312                let decoded = decoder.decode(buffered.as_ref())?;
313                if decoded > 0 {
314                    buffered.advance(decoded);
315                    continue;
316                }
317
318                if decoder.capacity() == 0 || input_finished {
319                    break;
320                }
321
322                if buffered.is_empty() {
323                    continue;
324                }
325
326                return Poll::Ready(Some(Err(ArrowError::ParseError(
327                    "CSV decoder made no progress while input bytes remain".to_string(),
328                ))));
329            }
330
331            match decoder.flush() {
332                Ok(Some(batch)) => return Poll::Ready(Some(Ok(batch))),
333                Ok(None) if input_finished => return Poll::Ready(None),
334                Ok(None) => continue,
335                Err(error) if is_skippable_arrow_error(&error) => {
336                    warn!(
337                        "Skipping bad CSV record while copying from {}: {}",
338                        path, error
339                    );
340                }
341                Err(error) => return Poll::Ready(Some(Err(error))),
342            }
343        }
344    })
345    .map(|result: std::result::Result<RecordBatch, ArrowError>| result.map_err(Into::into));
346
347    Ok(Box::pin(RecordBatchStreamAdapter::new(
348        projected_schema,
349        stream,
350    )))
351}
352
353pub fn is_skippable_arrow_error(error: &ArrowError) -> bool {
354    matches!(
355        error,
356        ArrowError::ParseError(_)
357            | ArrowError::CastError(_)
358            | ArrowError::ComputeError(_)
359            | ArrowError::InvalidArgumentError(_)
360    )
361}
362
363#[cfg(test)]
364mod tests {
365    use std::sync::Arc;
366
367    use arrow_schema::{DataType, Field};
368    use common_recordbatch::adapter::DfRecordBatchStreamAdapter;
369    use common_recordbatch::{RecordBatch, RecordBatches};
370    use common_test_util::find_workspace_path;
371    use datafusion::datasource::physical_plan::{CsvSource, FileSource};
372    use datatypes::prelude::ConcreteDataType;
373    use datatypes::schema::{ColumnSchema, Schema};
374    use datatypes::vectors::{Float64Vector, StringVector, UInt32Vector, VectorRef};
375    use futures::TryStreamExt;
376
377    use super::*;
378    use crate::file_format::{
379        FORMAT_COMPRESSION_TYPE, FORMAT_DELIMITER, FORMAT_HAS_HEADER, FORMAT_HEADERS,
380        FORMAT_SCHEMA_INFER_MAX_RECORD, FORMAT_SKIP_BAD_RECORDS, FORMAT_STRICT_HEADERS, FileFormat,
381        file_to_stream,
382    };
383    use crate::test_util::{format_schema, test_store};
384
385    fn test_data_root() -> String {
386        find_workspace_path("/src/common/datasource/tests/csv")
387            .display()
388            .to_string()
389    }
390
391    #[tokio::test]
392    async fn infer_schema_basic() {
393        let csv = CsvFormat::default();
394        let store = test_store(&test_data_root());
395        let schema = csv.infer_schema(&store, "simple.csv").await.unwrap();
396        let formatted: Vec<_> = format_schema(schema);
397
398        assert_eq!(
399            vec![
400                "c1: Utf8: NULL",
401                "c2: Int64: NULL",
402                "c3: Int64: NULL",
403                "c4: Int64: NULL",
404                "c5: Int64: NULL",
405                "c6: Int64: NULL",
406                "c7: Int64: NULL",
407                "c8: Int64: NULL",
408                "c9: Int64: NULL",
409                "c10: Utf8: NULL",
410                "c11: Float64: NULL",
411                "c12: Float64: NULL",
412                "c13: Utf8: NULL"
413            ],
414            formatted,
415        );
416    }
417
418    #[tokio::test]
419    async fn normalize_infer_schema() {
420        let csv = CsvFormat {
421            schema_infer_max_record: Some(3),
422            ..CsvFormat::default()
423        };
424        let store = test_store(&test_data_root());
425        let schema = csv.infer_schema(&store, "max_infer.csv").await.unwrap();
426        let formatted: Vec<_> = format_schema(schema);
427
428        assert_eq!(
429            vec![
430                "num: Int64: NULL",
431                "str: Utf8: NULL",
432                "ts: Utf8: NULL",
433                "t: Utf8: NULL",
434                "date: Date32: NULL"
435            ],
436            formatted,
437        );
438    }
439
440    #[tokio::test]
441    async fn infer_schema_with_limit() {
442        let csv = CsvFormat {
443            schema_infer_max_record: Some(3),
444            ..CsvFormat::default()
445        };
446        let store = test_store(&test_data_root());
447        let schema = csv
448            .infer_schema(&store, "schema_infer_limit.csv")
449            .await
450            .unwrap();
451        let formatted: Vec<_> = format_schema(schema);
452
453        assert_eq!(
454            vec![
455                "a: Int64: NULL",
456                "b: Float64: NULL",
457                "c: Int64: NULL",
458                "d: Int64: NULL"
459            ],
460            formatted
461        );
462
463        let csv = CsvFormat::default();
464        let store = test_store(&test_data_root());
465        let schema = csv
466            .infer_schema(&store, "schema_infer_limit.csv")
467            .await
468            .unwrap();
469        let formatted: Vec<_> = format_schema(schema);
470
471        assert_eq!(
472            vec![
473                "a: Int64: NULL",
474                "b: Float64: NULL",
475                "c: Int64: NULL",
476                "d: Utf8: NULL"
477            ],
478            formatted
479        );
480    }
481
482    #[test]
483    fn test_try_from() {
484        let map = HashMap::new();
485        let format: CsvFormat = CsvFormat::try_from(&map).unwrap();
486
487        assert_eq!(format, CsvFormat::default());
488
489        let map = HashMap::from([
490            (
491                FORMAT_SCHEMA_INFER_MAX_RECORD.to_string(),
492                "2000".to_string(),
493            ),
494            (FORMAT_COMPRESSION_TYPE.to_string(), "zstd".to_string()),
495            (FORMAT_DELIMITER.to_string(), b'\t'.to_string()),
496            (FORMAT_HAS_HEADER.to_string(), "false".to_string()),
497        ]);
498        let format = CsvFormat::try_from(&map).unwrap();
499
500        assert_eq!(
501            format,
502            CsvFormat {
503                compression_type: CompressionType::Zstd,
504                schema_infer_max_record: Some(2000),
505                delimiter: b'\t',
506                has_header: false,
507                skip_bad_records: false,
508                strict_headers: false,
509                timestamp_format: None,
510                time_format: None,
511                date_format: None
512            }
513        );
514
515        let map = HashMap::from([(FORMAT_SKIP_BAD_RECORDS.to_string(), "true".to_string())]);
516        let format = CsvFormat::try_from(&map).unwrap();
517
518        assert_eq!(
519            format,
520            CsvFormat {
521                skip_bad_records: true,
522                ..CsvFormat::default()
523            }
524        );
525
526        let map = HashMap::from([(FORMAT_STRICT_HEADERS.to_string(), "true".to_string())]);
527        let format = CsvFormat::try_from(&map).unwrap();
528
529        assert_eq!(
530            format,
531            CsvFormat {
532                strict_headers: true,
533                ..CsvFormat::default()
534            }
535        );
536
537        let map = HashMap::from([(FORMAT_HEADERS.to_string(), "true".to_string())]);
538        let format = CsvFormat::try_from(&map).unwrap();
539        assert_eq!(format, CsvFormat::default());
540
541        let map = HashMap::from([(FORMAT_HEADERS.to_string(), "false".to_string())]);
542        let format = CsvFormat::try_from(&map).unwrap();
543        assert_eq!(
544            format,
545            CsvFormat {
546                has_header: false,
547                ..CsvFormat::default()
548            }
549        );
550
551        let map = HashMap::from([
552            (FORMAT_HEADERS.to_string(), "false".to_string()),
553            (FORMAT_HAS_HEADER.to_string(), "false".to_string()),
554        ]);
555        let format = CsvFormat::try_from(&map).unwrap();
556        assert_eq!(
557            format,
558            CsvFormat {
559                has_header: false,
560                ..CsvFormat::default()
561            }
562        );
563    }
564
565    #[test]
566    fn test_try_from_rejects_invalid_bool_options() {
567        let map = HashMap::from([(FORMAT_SKIP_BAD_RECORDS.to_string(), "yes".to_string())]);
568        assert!(CsvFormat::try_from(&map).is_err());
569
570        let map = HashMap::from([(FORMAT_HEADERS.to_string(), "yes".to_string())]);
571        assert!(CsvFormat::try_from(&map).is_err());
572
573        let map = HashMap::from([(FORMAT_STRICT_HEADERS.to_string(), "yes".to_string())]);
574        assert!(CsvFormat::try_from(&map).is_err());
575    }
576
577    #[test]
578    fn test_try_from_rejects_conflicting_header_options() {
579        let map = HashMap::from([
580            (FORMAT_HEADERS.to_string(), "false".to_string()),
581            (FORMAT_HAS_HEADER.to_string(), "true".to_string()),
582        ]);
583        assert!(CsvFormat::try_from(&map).is_err());
584    }
585
586    #[test]
587    fn test_try_from_rejects_strict_headerless_csv() {
588        let map = HashMap::from([
589            (FORMAT_HEADERS.to_string(), "false".to_string()),
590            (FORMAT_STRICT_HEADERS.to_string(), "true".to_string()),
591        ]);
592        assert!(CsvFormat::try_from(&map).is_err());
593    }
594
595    #[tokio::test]
596    async fn test_compressed_csv() {
597        // Create test data
598        let column_schemas = vec![
599            ColumnSchema::new("id", ConcreteDataType::uint32_datatype(), false),
600            ColumnSchema::new("name", ConcreteDataType::string_datatype(), false),
601            ColumnSchema::new("value", ConcreteDataType::float64_datatype(), false),
602        ];
603        let schema = Arc::new(Schema::new(column_schemas));
604
605        // Create multiple record batches with different data
606        let batch1_columns: Vec<VectorRef> = vec![
607            Arc::new(UInt32Vector::from_slice(vec![1, 2, 3])),
608            Arc::new(StringVector::from(vec!["Alice", "Bob", "Charlie"])),
609            Arc::new(Float64Vector::from_slice(vec![10.5, 20.3, 30.7])),
610        ];
611        let batch1 = RecordBatch::new(schema.clone(), batch1_columns).unwrap();
612
613        let batch2_columns: Vec<VectorRef> = vec![
614            Arc::new(UInt32Vector::from_slice(vec![4, 5, 6])),
615            Arc::new(StringVector::from(vec!["David", "Eva", "Frank"])),
616            Arc::new(Float64Vector::from_slice(vec![40.1, 50.2, 60.3])),
617        ];
618        let batch2 = RecordBatch::new(schema.clone(), batch2_columns).unwrap();
619
620        let batch3_columns: Vec<VectorRef> = vec![
621            Arc::new(UInt32Vector::from_slice(vec![7, 8, 9])),
622            Arc::new(StringVector::from(vec!["Grace", "Henry", "Ivy"])),
623            Arc::new(Float64Vector::from_slice(vec![70.4, 80.5, 90.6])),
624        ];
625        let batch3 = RecordBatch::new(schema.clone(), batch3_columns).unwrap();
626
627        // Combine all batches into a RecordBatches collection
628        let recordbatches = RecordBatches::try_new(schema, vec![batch1, batch2, batch3]).unwrap();
629
630        // Test with different compression types
631        let compression_types = vec![
632            CompressionType::Gzip,
633            CompressionType::Bzip2,
634            CompressionType::Xz,
635            CompressionType::Zstd,
636        ];
637
638        // Create a temporary file path
639        let temp_dir = common_test_util::temp_dir::create_temp_dir("test_compressed_csv");
640        for compression_type in compression_types {
641            let format = CsvFormat {
642                compression_type,
643                ..CsvFormat::default()
644            };
645
646            // Use correct format without Debug formatter
647            let compressed_file_name =
648                format!("test_compressed_csv.{}", compression_type.file_extension());
649            let compressed_file_path = temp_dir.path().join(&compressed_file_name);
650            let compressed_file_path_str = compressed_file_path.to_str().unwrap();
651
652            // Create a simple file store for testing
653            let store = test_store("/");
654
655            // Export CSV with compression
656            let rows = stream_to_csv(
657                Box::pin(DfRecordBatchStreamAdapter::new(recordbatches.as_stream())),
658                store,
659                compressed_file_path_str,
660                1024,
661                1,
662                &format,
663            )
664            .await
665            .unwrap();
666
667            assert_eq!(rows, 9);
668
669            // Verify compressed file was created and has content
670            assert!(compressed_file_path.exists());
671            let file_size = std::fs::metadata(&compressed_file_path).unwrap().len();
672            assert!(file_size > 0);
673
674            // Verify the file is actually compressed
675            let file_content = std::fs::read(&compressed_file_path).unwrap();
676            // Compressed files should not start with CSV header
677            // They should have compression magic bytes
678            match compression_type {
679                CompressionType::Gzip => {
680                    // Gzip magic bytes: 0x1f 0x8b
681                    assert_eq!(file_content[0], 0x1f, "Gzip file should start with 0x1f");
682                    assert_eq!(
683                        file_content[1], 0x8b,
684                        "Gzip file should have 0x8b as second byte"
685                    );
686                }
687                CompressionType::Bzip2 => {
688                    // Bzip2 magic bytes: 'BZ'
689                    assert_eq!(file_content[0], b'B', "Bzip2 file should start with 'B'");
690                    assert_eq!(
691                        file_content[1], b'Z',
692                        "Bzip2 file should have 'Z' as second byte"
693                    );
694                }
695                CompressionType::Xz => {
696                    // XZ magic bytes: 0xFD '7zXZ'
697                    assert_eq!(file_content[0], 0xFD, "XZ file should start with 0xFD");
698                }
699                CompressionType::Zstd => {
700                    // Zstd magic bytes: 0x28 0xB5 0x2F 0xFD
701                    assert_eq!(file_content[0], 0x28, "Zstd file should start with 0x28");
702                    assert_eq!(
703                        file_content[1], 0xB5,
704                        "Zstd file should have 0xB5 as second byte"
705                    );
706                }
707                _ => {}
708            }
709
710            // Verify the compressed file can be decompressed and content matches original data
711            let store = test_store("/");
712            let schema = Arc::new(
713                CsvFormat {
714                    compression_type,
715                    ..Default::default()
716                }
717                .infer_schema(&store, compressed_file_path_str)
718                .await
719                .unwrap(),
720            );
721            let csv_source = CsvSource::new(schema).with_batch_size(8192);
722
723            let stream = file_to_stream(
724                &store,
725                compressed_file_path_str,
726                csv_source.clone(),
727                None,
728                compression_type,
729            )
730            .await
731            .unwrap();
732
733            let batches = stream.try_collect::<Vec<_>>().await.unwrap();
734            let pretty_print = arrow::util::pretty::pretty_format_batches(&batches)
735                .unwrap()
736                .to_string();
737            let expected = r#"+----+---------+-------+
738| id | name    | value |
739+----+---------+-------+
740| 1  | Alice   | 10.5  |
741| 2  | Bob     | 20.3  |
742| 3  | Charlie | 30.7  |
743| 4  | David   | 40.1  |
744| 5  | Eva     | 50.2  |
745| 6  | Frank   | 60.3  |
746| 7  | Grace   | 70.4  |
747| 8  | Henry   | 80.5  |
748| 9  | Ivy     | 90.6  |
749+----+---------+-------+"#;
750            assert_eq!(expected, pretty_print);
751        }
752    }
753
754    #[tokio::test]
755    async fn test_tolerant_csv_stream_continues_after_parse_error() {
756        let temp_dir = common_test_util::temp_dir::create_temp_dir("test_tolerant_csv_stream");
757        let csv_file_path = temp_dir.path().join("input.csv");
758        std::fs::write(
759            &csv_file_path,
760            "id,name,value\n1,Alice,10.5\nbad,Bad,20.0\nworse,Bad,21.0\n2,Bob,30.5",
761        )
762        .unwrap();
763
764        let store = test_store("/");
765        let schema = Arc::new(arrow_schema::Schema::new(vec![
766            Field::new("id", DataType::UInt32, false),
767            Field::new("name", DataType::Utf8, false),
768            Field::new("value", DataType::Float64, false),
769        ]));
770        let path = csv_file_path.to_str().unwrap();
771
772        let stream =
773            tolerant_csv_stream(&store, path, schema, vec![0, 1, 2], &CsvFormat::default())
774                .await
775                .unwrap();
776        let batches = stream.try_collect::<Vec<_>>().await.unwrap();
777        let pretty_print = arrow::util::pretty::pretty_format_batches(&batches)
778            .unwrap()
779            .to_string();
780        let expected = r#"+----+-------+-------+
781| id | name  | value |
782+----+-------+-------+
783| 1  | Alice | 10.5  |
784| 2  | Bob   | 30.5  |
785+----+-------+-------+"#;
786        assert_eq!(expected, pretty_print);
787    }
788
789    #[tokio::test]
790    async fn test_tolerant_csv_stream_fails_on_structural_csv_error() {
791        let temp_dir =
792            common_test_util::temp_dir::create_temp_dir("test_tolerant_csv_stream_csv_error");
793        let csv_file_path = temp_dir.path().join("input.csv");
794        std::fs::write(&csv_file_path, "id,name,value\n1,Alice,10.5\n2,Bob\n").unwrap();
795
796        let store = test_store("/");
797        let schema = Arc::new(arrow_schema::Schema::new(vec![
798            Field::new("id", DataType::UInt32, false),
799            Field::new("name", DataType::Utf8, false),
800            Field::new("value", DataType::Float64, false),
801        ]));
802        let path = csv_file_path.to_str().unwrap();
803
804        let stream =
805            tolerant_csv_stream(&store, path, schema, vec![0, 1, 2], &CsvFormat::default())
806                .await
807                .unwrap();
808        let error = stream.try_collect::<Vec<_>>().await.unwrap_err();
809
810        assert!(error.to_string().contains("incorrect number of fields"));
811    }
812}
common_datasource/file_format/csv.rs

common_datasource/file_format/
csv.rs
