ddb_core/file_io/
csv_reader.rs

1// CSV/Delimited file reader with memory-mapped I/O support for large files
2use crate::engine::row::Row;
3use crate::error::{DdbError, Result};
4use crate::functions::Value;
5use memmap2::Mmap;
6use std::fs::File;
7use std::io::{BufRead, BufReader};
8use std::path::Path;
9
10/// Threshold for using memory-mapped I/O (10MB)
11const MMAP_THRESHOLD: u64 = 10 * 1024 * 1024;
12
13enum ReaderBackend {
14    Buffered(BufReader<File>),
15    MemoryMapped {
16        mmap: Mmap,
17        position: usize,
18    },
19}
20
21pub struct CsvReader {
22    backend: ReaderBackend,
23    delimiter: char,
24    column_names: Vec<String>,
25    line_number: usize,
26    has_header: bool,
27}
28
29impl CsvReader {
30    /// Create a new CSV reader with specified delimiter
31    /// Automatically uses memory-mapped I/O for files larger than 10MB
32    pub fn new<P: AsRef<Path>>(path: P, delimiter: char, has_header: bool) -> Result<Self> {
33        let path_ref = path.as_ref();
34        let file = File::open(path_ref)?;
35        let metadata = file.metadata()?;
36        let file_size = metadata.len();
37
38        let backend = if file_size >= MMAP_THRESHOLD {
39            // Use memory-mapped I/O for large files
40            let mmap = unsafe { Mmap::map(&file)? };
41            ReaderBackend::MemoryMapped {
42                mmap,
43                position: 0,
44            }
45        } else {
46            // Use buffered I/O for small files
47            ReaderBackend::Buffered(BufReader::new(file))
48        };
49
50        Ok(Self {
51            backend,
52            delimiter,
53            column_names: Vec::new(),
54            line_number: 0,
55            has_header,
56        })
57    }
58
59    /// Create a CSV reader with explicit column names (no header in file)
60    pub fn with_columns<P: AsRef<Path>>(
61        path: P,
62        delimiter: char,
63        columns: Vec<String>,
64    ) -> Result<Self> {
65        let mut reader = Self::new(path, delimiter, false)?;
66        reader.column_names = columns;
67        Ok(reader)
68    }
69
70    /// Get the column names
71    pub fn column_names(&self) -> &[String] {
72        &self.column_names
73    }
74
75    /// Read the next row
76    pub fn next_row(&mut self) -> Result<Option<Row>> {
77        let line_opt = match &mut self.backend {
78            ReaderBackend::Buffered(reader) => {
79                let mut line = String::new();
80                match reader.read_line(&mut line) {
81                    Ok(0) => None,
82                    Ok(_) => Some(line.trim_end_matches(&['\r', '\n'][..]).to_string()),
83                    Err(e) => return Err(DdbError::IoError(e)),
84                }
85            }
86            ReaderBackend::MemoryMapped { mmap, position } => {
87                Self::read_line_from_mmap(mmap, position)?
88            }
89        };
90
91        match line_opt {
92            None => Ok(None), // EOF
93            Some(line) => {
94                self.line_number += 1;
95
96                // If this is the first line and we have a header, parse column names
97                if self.line_number == 1 && self.has_header {
98                    self.column_names = self.parse_line(&line);
99                    return self.next_row(); // Read next line for actual data
100                }
101
102                // If we don't have column names, generate them
103                if self.column_names.is_empty() {
104                    let values = self.parse_line(&line);
105                    self.column_names = (0..values.len())
106                        .map(|i| format!("col{}", i))
107                        .collect();
108                }
109
110                let values = self.parse_line(&line);
111                let row = self.create_row(values)?;
112                Ok(Some(row))
113            }
114        }
115    }
116
117    /// Read a line from memory-mapped file
118    fn read_line_from_mmap(mmap: &Mmap, position: &mut usize) -> Result<Option<String>> {
119        if *position >= mmap.len() {
120            return Ok(None);
121        }
122
123        // Find the next newline
124        let start = *position;
125        let mut end = start;
126
127        while end < mmap.len() && mmap[end] != b'\n' {
128            end += 1;
129        }
130
131        // Extract the line
132        let line_bytes = &mmap[start..end];
133        let mut line = String::from_utf8_lossy(line_bytes).to_string();
134
135        // Trim \r if present (for Windows line endings)
136        if line.ends_with('\r') {
137            line.pop();
138        }
139
140        // Update position to after the newline
141        *position = if end < mmap.len() { end + 1 } else { end };
142
143        Ok(Some(line))
144    }
145
146    /// Parse a line into fields
147    fn parse_line(&self, line: &str) -> Vec<String> {
148        let mut fields = Vec::new();
149        let mut current_field = String::new();
150        let mut in_quotes = false;
151        let mut chars = line.chars().peekable();
152
153        while let Some(ch) = chars.next() {
154            if ch == '"' {
155                in_quotes = !in_quotes;
156            } else if ch == self.delimiter && !in_quotes {
157                fields.push(current_field.clone());
158                current_field.clear();
159            } else {
160                current_field.push(ch);
161            }
162        }
163
164        // Add the last field
165        fields.push(current_field);
166
167        fields
168    }
169
170    /// Create a Row from parsed field values
171    fn create_row(&self, fields: Vec<String>) -> Result<Row> {
172        if fields.len() != self.column_names.len() {
173            return Err(DdbError::InvalidColumnCount {
174                expected: self.column_names.len(),
175                actual: fields.len(),
176            });
177        }
178
179        let mut row = Row::new();
180        for (name, field) in self.column_names.iter().zip(fields.iter()) {
181            let value = Value::from_str(field);
182            row.set(name, value);
183        }
184
185        Ok(row)
186    }
187
188    /// Get current line number
189    pub fn line_number(&self) -> usize {
190        self.line_number
191    }
192}
193
194#[cfg(test)]
195mod tests {
196    use super::*;
197    use std::io::Write;
198    use tempfile::NamedTempFile;
199
200    #[test]
201    fn test_csv_reader_with_header() {
202        // Create a temporary CSV file
203        let mut file = NamedTempFile::new().unwrap();
204        writeln!(file, "id,name,age").unwrap();
205        writeln!(file, "1,Alice,30").unwrap();
206        writeln!(file, "2,Bob,25").unwrap();
207        file.flush().unwrap();
208
209        let mut reader = CsvReader::new(file.path(), ',', true).unwrap();
210
211        // Read first row
212        let row1 = reader.next_row().unwrap().unwrap();
213        assert_eq!(row1.get("id"), Some(&Value::Integer(1)));
214        assert_eq!(row1.get("name"), Some(&Value::String("Alice".to_string())));
215        assert_eq!(row1.get("age"), Some(&Value::Integer(30)));
216
217        // Read second row
218        let row2 = reader.next_row().unwrap().unwrap();
219        assert_eq!(row2.get("id"), Some(&Value::Integer(2)));
220
221        // EOF
222        assert!(reader.next_row().unwrap().is_none());
223    }
224
225    #[test]
226    fn test_csv_reader_without_header() {
227        let mut file = NamedTempFile::new().unwrap();
228        writeln!(file, "1,Alice,30").unwrap();
229        writeln!(file, "2,Bob,25").unwrap();
230        file.flush().unwrap();
231
232        let columns = vec!["id".to_string(), "name".to_string(), "age".to_string()];
233        let mut reader = CsvReader::with_columns(file.path(), ',', columns).unwrap();
234
235        let row = reader.next_row().unwrap().unwrap();
236        assert_eq!(row.get("id"), Some(&Value::Integer(1)));
237        assert_eq!(row.get("name"), Some(&Value::String("Alice".to_string())));
238    }
239
240    #[test]
241    fn test_csv_reader_quoted_fields() {
242        let mut file = NamedTempFile::new().unwrap();
243        writeln!(file, "name,description").unwrap();
244        writeln!(file, r#"Test,"A field, with comma""#).unwrap();
245        file.flush().unwrap();
246
247        let mut reader = CsvReader::new(file.path(), ',', true).unwrap();
248        let row = reader.next_row().unwrap().unwrap();
249
250        assert_eq!(row.get("name"), Some(&Value::String("Test".to_string())));
251        assert_eq!(row.get("description"), Some(&Value::String("A field, with comma".to_string())));
252    }
253}
ddb_core/file_io/csv_reader.rs

ddb_core/file_io/
csv_reader.rs
