sql_cli/data/
datatable_loaders.rs

1use crate::data::stream_loader::{detect_delimiter_from_path, strip_field_padding, CsvReadOptions};
2use crate::datatable::{DataColumn, DataRow, DataTable, DataType, DataValue};
3use anyhow::{Context, Result};
4use csv::ReaderBuilder;
5use serde_json::Value as JsonValue;
6use std::collections::HashSet;
7use std::fs::File;
8use std::io::{BufRead, BufReader, Read};
9use std::path::Path;
10
11/// Helper to detect if a field in the raw CSV line is a null (unquoted empty).
12/// `delimiter` is the field separator character used in the source.
13fn is_null_field(raw_line: &str, field_index: usize, delimiter: char) -> bool {
14    let mut delim_count = 0;
15    let mut in_quotes = false;
16    let mut field_start = 0;
17    let mut prev_char = ' ';
18
19    for (i, ch) in raw_line.char_indices() {
20        if ch == '"' && prev_char != '\\' {
21            in_quotes = !in_quotes;
22        }
23
24        if ch == delimiter && !in_quotes {
25            if delim_count == field_index {
26                let field_end = i;
27                let field_content = &raw_line[field_start..field_end].trim();
28                // If empty, check if it was quoted (quoted empty = empty string, unquoted empty = NULL)
29                if field_content.is_empty() {
30                    return true; // Unquoted empty field -> NULL
31                }
32                // If it starts and ends with quotes but is empty inside, it's an empty string, not NULL
33                if field_content.starts_with('"')
34                    && field_content.ends_with('"')
35                    && field_content.len() == 2
36                {
37                    return false; // Quoted empty field -> empty string
38                }
39                return false; // Non-empty field -> not NULL
40            }
41            delim_count += 1;
42            field_start = i + ch.len_utf8();
43        }
44        prev_char = ch;
45    }
46
47    // Check last field
48    if delim_count == field_index {
49        let field_content = raw_line[field_start..]
50            .trim()
51            .trim_end_matches('\n')
52            .trim_end_matches('\r');
53        // If empty, check if it was quoted
54        if field_content.is_empty() {
55            return true; // Unquoted empty field -> NULL
56        }
57        // If it starts and ends with quotes but is empty inside, it's an empty string, not NULL
58        if field_content.starts_with('"')
59            && field_content.ends_with('"')
60            && field_content.len() == 2
61        {
62            return false; // Quoted empty field -> empty string
63        }
64        return false; // Non-empty field -> not NULL
65    }
66
67    false // Field not found -> not NULL (shouldn't happen)
68}
69
70/// Load a CSV file into a `DataTable`. Delimiter is auto-detected from the
71/// file extension (`.tsv` → tab, `.psv` → pipe, else comma). To override the
72/// auto-detect, use [`load_csv_to_datatable_with_opts`].
73pub fn load_csv_to_datatable<P: AsRef<Path>>(path: P, table_name: &str) -> Result<DataTable> {
74    let path_ref = path.as_ref();
75    let opts = CsvReadOptions {
76        delimiter: detect_delimiter_from_path(&path_ref.display().to_string()),
77        has_headers: true,
78    };
79    load_csv_to_datatable_with_opts(path, table_name, &opts)
80}
81
82/// Load a CSV file into a `DataTable` honouring caller-supplied options
83/// (delimiter, headers).
84pub fn load_csv_to_datatable_with_opts<P: AsRef<Path>>(
85    path: P,
86    table_name: &str,
87    opts: &CsvReadOptions,
88) -> Result<DataTable> {
89    let mut file = File::open(&path)
90        .with_context(|| format!("Failed to open CSV file: {:?}", path.as_ref()))?;
91    let mut raw_buffer = Vec::new();
92    file.read_to_end(&mut raw_buffer)
93        .with_context(|| format!("Failed to read CSV file: {:?}", path.as_ref()))?;
94
95    // Strip alignment padding up front so both passes below agree on where
96    // fields start and end (the NULL pass indexes into these same bytes).
97    let buffer = strip_field_padding(&raw_buffer, opts.delimiter);
98
99    let mut reader = ReaderBuilder::new()
100        .has_headers(opts.has_headers)
101        .delimiter(opts.delimiter)
102        .from_reader(&buffer[..]);
103
104    // Get headers and create columns
105    let headers = reader.headers()?.clone();
106    let mut table = DataTable::new(table_name);
107
108    // Add metadata about the source
109    table
110        .metadata
111        .insert("source_type".to_string(), "csv".to_string());
112    table.metadata.insert(
113        "source_path".to_string(),
114        path.as_ref().display().to_string(),
115    );
116    table.metadata.insert(
117        "delimiter".to_string(),
118        match opts.delimiter {
119            b'\t' => "\\t".to_string(),
120            b'\n' => "\\n".to_string(),
121            b'\r' => "\\r".to_string(),
122            b => (b as char).to_string(),
123        },
124    );
125
126    // Create columns from headers (types will be inferred later)
127    for header in &headers {
128        table.add_column(DataColumn::new(header));
129    }
130
131    // Second view over the same (padding-stripped) bytes for raw line reading
132    let mut line_reader = BufReader::new(&buffer[..]);
133    let mut raw_line = String::new();
134    // Skip header line
135    line_reader.read_line(&mut raw_line)?;
136
137    // Read all rows first to collect data
138    let mut string_rows = Vec::new();
139    let mut raw_lines = Vec::new();
140
141    for result in reader.records() {
142        let record = result?;
143        let row: Vec<String> = record
144            .iter()
145            .map(std::string::ToString::to_string)
146            .collect();
147
148        // Read the corresponding raw line
149        raw_line.clear();
150        line_reader.read_line(&mut raw_line)?;
151        raw_lines.push(raw_line.clone());
152
153        string_rows.push(row);
154    }
155
156    // Infer column types by sampling the data
157    let mut column_types = vec![DataType::Null; headers.len()];
158    let sample_size = string_rows.len().min(100); // Sample first 100 rows for type inference
159
160    for row in string_rows.iter().take(sample_size) {
161        for (col_idx, value) in row.iter().enumerate() {
162            if !value.is_empty() {
163                let inferred = DataType::infer_from_string(value);
164                column_types[col_idx] = column_types[col_idx].merge(&inferred);
165            }
166        }
167    }
168
169    // Update column types
170    for (col_idx, column) in table.columns.iter_mut().enumerate() {
171        column.data_type = column_types[col_idx].clone();
172    }
173
174    // Convert string data to typed DataValues and add rows
175    for (row_idx, string_row) in string_rows.iter().enumerate() {
176        let mut values = Vec::new();
177        let raw_line = &raw_lines[row_idx];
178
179        for (col_idx, value) in string_row.iter().enumerate() {
180            let data_value = if value.is_empty() {
181                // Distinguish between NULL (,,) and empty string ("")
182                if is_null_field(raw_line, col_idx, opts.delimiter as char) {
183                    DataValue::Null
184                } else {
185                    DataValue::String(String::new())
186                }
187            } else {
188                DataValue::from_string(value, &column_types[col_idx])
189            };
190            values.push(data_value);
191        }
192        table
193            .add_row(DataRow::new(values))
194            .map_err(|e| anyhow::anyhow!(e))?;
195    }
196
197    // Update column statistics
198    table.infer_column_types();
199
200    Ok(table)
201}
202
203/// Load a JSON file into a `DataTable`.
204///
205/// Accepts either a JSON array of objects (`[{...}, {...}]`) or JSONL
206/// (one JSON object per line). Format is auto-detected.
207pub fn load_json_to_datatable<P: AsRef<Path>>(path: P, table_name: &str) -> Result<DataTable> {
208    // Read file as string first to preserve key order
209    let mut file = File::open(&path)
210        .with_context(|| format!("Failed to open JSON file: {:?}", path.as_ref()))?;
211    let mut json_str = String::new();
212    file.read_to_string(&mut json_str)?;
213
214    let json_data: Vec<JsonValue> = crate::data::stream_loader::parse_json_records(&json_str)?;
215
216    if json_data.is_empty() {
217        return Ok(DataTable::new(table_name));
218    }
219
220    // Schema is the union of keys across the first 100 records so heterogeneous
221    // JSONL streams don't silently drop columns missing on the first object.
222    let column_names = crate::data::stream_loader::collect_column_names(&json_data, 100);
223    if column_names.is_empty() {
224        return Err(anyhow::anyhow!(
225            "JSON data must contain objects (got non-object records)"
226        ));
227    }
228
229    let mut table = DataTable::new(table_name);
230
231    // Add metadata
232    table
233        .metadata
234        .insert("source_type".to_string(), "json".to_string());
235    table.metadata.insert(
236        "source_path".to_string(),
237        path.as_ref().display().to_string(),
238    );
239
240    for name in &column_names {
241        table.add_column(DataColumn::new(name));
242    }
243
244    // Collect all values as strings first for type inference
245    let mut string_rows = Vec::new();
246    for json_obj in &json_data {
247        if let Some(obj) = json_obj.as_object() {
248            let mut row = Vec::new();
249            for name in &column_names {
250                let value_str = match obj.get(name) {
251                    Some(JsonValue::Null) | None => String::new(),
252                    Some(JsonValue::Bool(b)) => b.to_string(),
253                    Some(JsonValue::Number(n)) => n.to_string(),
254                    Some(JsonValue::String(s)) => s.clone(),
255                    Some(JsonValue::Array(arr)) => format!("{arr:?}"), // Arrays as debug string for now
256                    Some(JsonValue::Object(obj)) => format!("{obj:?}"), // Objects as debug string for now
257                };
258                row.push(value_str);
259            }
260            string_rows.push(row);
261        }
262    }
263
264    // Infer column types
265    let mut column_types = vec![DataType::Null; column_names.len()];
266    let sample_size = string_rows.len().min(100);
267
268    for row in string_rows.iter().take(sample_size) {
269        for (col_idx, value) in row.iter().enumerate() {
270            if !value.is_empty() {
271                let inferred = DataType::infer_from_string(value);
272                column_types[col_idx] = column_types[col_idx].merge(&inferred);
273            }
274        }
275    }
276
277    // Update column types
278    for (col_idx, column) in table.columns.iter_mut().enumerate() {
279        column.data_type = column_types[col_idx].clone();
280    }
281
282    // Convert to DataRows
283    for string_row in string_rows {
284        let mut values = Vec::new();
285        for (col_idx, value) in string_row.iter().enumerate() {
286            let data_value = DataValue::from_string(value, &column_types[col_idx]);
287            values.push(data_value);
288        }
289        table
290            .add_row(DataRow::new(values))
291            .map_err(|e| anyhow::anyhow!(e))?;
292    }
293
294    // Update statistics
295    table.infer_column_types();
296
297    Ok(table)
298}
299
300/// Load JSON data directly (already parsed) into a `DataTable`
301pub fn load_json_data_to_datatable(data: Vec<JsonValue>, table_name: &str) -> Result<DataTable> {
302    if data.is_empty() {
303        return Ok(DataTable::new(table_name));
304    }
305
306    // Extract column names from all objects (union of all keys)
307    let mut all_columns = HashSet::new();
308    for item in &data {
309        if let Some(obj) = item.as_object() {
310            for key in obj.keys() {
311                all_columns.insert(key.clone());
312            }
313        }
314    }
315
316    let column_names: Vec<String> = all_columns.into_iter().collect();
317    let mut table = DataTable::new(table_name);
318
319    // Add metadata
320    table
321        .metadata
322        .insert("source_type".to_string(), "json_data".to_string());
323
324    // Create columns
325    for name in &column_names {
326        table.add_column(DataColumn::new(name));
327    }
328
329    // Process data similar to file loading
330    let mut string_rows = Vec::new();
331    for json_obj in &data {
332        if let Some(obj) = json_obj.as_object() {
333            let mut row = Vec::new();
334            for name in &column_names {
335                let value_str = match obj.get(name) {
336                    Some(JsonValue::Null) | None => String::new(),
337                    Some(JsonValue::Bool(b)) => b.to_string(),
338                    Some(JsonValue::Number(n)) => n.to_string(),
339                    Some(JsonValue::String(s)) => s.clone(),
340                    Some(JsonValue::Array(arr)) => format!("{arr:?}"),
341                    Some(JsonValue::Object(obj)) => format!("{obj:?}"),
342                };
343                row.push(value_str);
344            }
345            string_rows.push(row);
346        }
347    }
348
349    // Infer types and convert to DataRows (same as above)
350    let mut column_types = vec![DataType::Null; column_names.len()];
351    let sample_size = string_rows.len().min(100);
352
353    for row in string_rows.iter().take(sample_size) {
354        for (col_idx, value) in row.iter().enumerate() {
355            if !value.is_empty() {
356                let inferred = DataType::infer_from_string(value);
357                column_types[col_idx] = column_types[col_idx].merge(&inferred);
358            }
359        }
360    }
361
362    for (col_idx, column) in table.columns.iter_mut().enumerate() {
363        column.data_type = column_types[col_idx].clone();
364    }
365
366    for string_row in string_rows {
367        let mut values = Vec::new();
368        for (col_idx, value) in string_row.iter().enumerate() {
369            let data_value = DataValue::from_string(value, &column_types[col_idx]);
370            values.push(data_value);
371        }
372        table
373            .add_row(DataRow::new(values))
374            .map_err(|e| anyhow::anyhow!(e))?;
375    }
376
377    table.infer_column_types();
378
379    Ok(table)
380}
381
382#[cfg(test)]
383mod tests {
384    use super::*;
385    use std::io::Write;
386    use tempfile::NamedTempFile;
387
388    #[test]
389    fn test_load_csv() -> Result<()> {
390        // Create a temporary CSV file
391        let mut temp_file = NamedTempFile::new()?;
392        writeln!(temp_file, "id,name,price,quantity")?;
393        writeln!(temp_file, "1,Widget,9.99,100")?;
394        writeln!(temp_file, "2,Gadget,19.99,50")?;
395        writeln!(temp_file, "3,Doohickey,5.00,200")?;
396        temp_file.flush()?;
397
398        let table = load_csv_to_datatable(temp_file.path(), "products")?;
399
400        assert_eq!(table.name, "products");
401        assert_eq!(table.column_count(), 4);
402        assert_eq!(table.row_count(), 3);
403
404        // Check column types were inferred correctly
405        assert_eq!(table.columns[0].name, "id");
406        assert_eq!(table.columns[0].data_type, DataType::Integer);
407
408        assert_eq!(table.columns[1].name, "name");
409        assert_eq!(table.columns[1].data_type, DataType::String);
410
411        assert_eq!(table.columns[2].name, "price");
412        assert_eq!(table.columns[2].data_type, DataType::Float);
413
414        assert_eq!(table.columns[3].name, "quantity");
415        assert_eq!(table.columns[3].data_type, DataType::Integer);
416
417        // Check data
418        let value = table.get_value_by_name(0, "name").unwrap();
419        assert_eq!(value.to_string(), "Widget");
420
421        Ok(())
422    }
423
424    #[test]
425    fn test_fractional_value_beyond_sample_window_promotes_to_float() -> Result<()> {
426        // Regression: type inference only samples the first 100 rows. A column
427        // that is all integers in the sample but has a fractional value further
428        // down used to demote that value to a String, which then sorted after
429        // every numeric value (String > Integer). It must be a Float instead.
430        let mut temp_file = NamedTempFile::new()?;
431        writeln!(temp_file, "id,area")?;
432        for i in 0..120 {
433            writeln!(temp_file, "{i},{}", i * 10)?; // all integers in the sample
434        }
435        writeln!(temp_file, "999,34.2")?; // fractional value past row 100
436        temp_file.flush()?;
437
438        let table = load_csv_to_datatable(temp_file.path(), "areas")?;
439
440        // The column re-merges to Float once the fractional value is seen.
441        let area_idx = table.get_column_index("area").unwrap();
442        assert_eq!(table.columns[area_idx].data_type, DataType::Float);
443
444        // The fractional value is stored as a number, not a String.
445        let last = table.get_value(120, area_idx).unwrap();
446        assert!(
447            matches!(last, DataValue::Float(f) if (*f - 34.2).abs() < 1e-9),
448            "expected Float(34.2), got {last:?}"
449        );
450
451        Ok(())
452    }
453
454    #[test]
455    fn test_load_json() -> Result<()> {
456        // Create a temporary JSON file
457        let mut temp_file = NamedTempFile::new()?;
458        writeln!(
459            temp_file,
460            r#"[
461            {{"id": 1, "name": "Alice", "active": true, "score": 95.5}},
462            {{"id": 2, "name": "Bob", "active": false, "score": 87.3}},
463            {{"id": 3, "name": "Charlie", "active": true, "score": null}}
464        ]"#
465        )?;
466        temp_file.flush()?;
467
468        let table = load_json_to_datatable(temp_file.path(), "users")?;
469
470        assert_eq!(table.name, "users");
471        assert_eq!(table.column_count(), 4);
472        assert_eq!(table.row_count(), 3);
473
474        // Check that null handling works
475        let score = table.get_value_by_name(2, "score").unwrap();
476        assert!(score.is_null());
477
478        Ok(())
479    }
480
481    #[test]
482    fn test_load_csv_with_pipe_delimiter_via_opts() -> Result<()> {
483        let mut temp_file = NamedTempFile::new()?;
484        writeln!(temp_file, "id|name|price")?;
485        writeln!(temp_file, "1|Widget|9.99")?;
486        writeln!(temp_file, "2|Gadget|19.99")?;
487        temp_file.flush()?;
488
489        let opts = CsvReadOptions {
490            delimiter: b'|',
491            has_headers: true,
492        };
493        let table = load_csv_to_datatable_with_opts(temp_file.path(), "psv_products", &opts)?;
494
495        assert_eq!(table.column_count(), 3);
496        assert_eq!(table.row_count(), 2);
497        assert_eq!(table.columns[0].name, "id");
498        assert_eq!(table.columns[1].name, "name");
499        assert_eq!(table.columns[0].data_type, DataType::Integer);
500        assert_eq!(
501            table.get_value_by_name(0, "name").unwrap().to_string(),
502            "Widget"
503        );
504        assert_eq!(
505            table.metadata.get("delimiter").map(String::as_str),
506            Some("|")
507        );
508        Ok(())
509    }
510
511    #[test]
512    fn test_default_load_csv_records_comma_delimiter() -> Result<()> {
513        let mut temp_file = NamedTempFile::new()?;
514        writeln!(temp_file, "a,b")?;
515        writeln!(temp_file, "1,2")?;
516        temp_file.flush()?;
517
518        let table = load_csv_to_datatable(temp_file.path(), "t")?;
519        assert_eq!(
520            table.metadata.get("delimiter").map(String::as_str),
521            Some(",")
522        );
523        Ok(())
524    }
525}
sql_cli/data/datatable_loaders.rs

sql_cli/data/
datatable_loaders.rs
