Never expect nil for csv io dtypes (#672)

elixir-explorer · Jul 28, 2023 · 74bdb47 · 74bdb47
1 parent bf9e5de
commit 74bdb47
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 24 deletions.
diff --git a/native/explorer/src/dataframe/io.rs b/native/explorer/src/dataframe/io.rs
@@ -49,7 +49,7 @@ pub fn df_from_csv(
     delimiter_as_byte: u8,
     do_rechunk: bool,
     column_names: Option<Vec<String>>,
-    dtypes: Option<Vec<(&str, &str)>>,
+    dtypes: Vec<(&str, &str)>,
     encoding: &str,
     null_vals: Vec<String>,
     parse_dates: bool,
@@ -60,12 +60,6 @@ pub fn df_from_csv(
         _ => CsvEncoding::Utf8,
     };
 
-    let schema = match dtypes {
-        Some(dtypes) => Some(schema_from_dtypes_pairs(dtypes)?),
-
-        None => None,
-    };
-
     let reader = CsvReader::from_path(filename)?
         .infer_schema(infer_schema_length)
         .has_header(has_header)
@@ -77,7 +71,7 @@ pub fn df_from_csv(
         .with_rechunk(do_rechunk)
         .with_encoding(encoding)
         .with_columns(column_names)
-        .with_dtypes(schema)
+        .with_dtypes(Some(schema_from_dtypes_pairs(dtypes)?))
         .with_null_values(Some(NullValues::AllColumns(null_vals)))
         .with_end_of_line_char(eol_delimiter.unwrap_or(b'\n'));
 
@@ -174,7 +168,7 @@ pub fn df_load_csv(
     delimiter_as_byte: u8,
     do_rechunk: bool,
     column_names: Option<Vec<String>>,
-    dtypes: Option<Vec<(&str, &str)>>,
+    dtypes: Vec<(&str, &str)>,
     encoding: &str,
     null_vals: Vec<String>,
     parse_dates: bool,
@@ -185,12 +179,6 @@ pub fn df_load_csv(
         _ => CsvEncoding::Utf8,
     };
 
-    let schema = match dtypes {
-        Some(dtypes) => Some(schema_from_dtypes_pairs(dtypes)?),
-
-        None => None,
-    };
-
     let cursor = Cursor::new(binary.as_slice());
 
     let reader = CsvReader::new(cursor)
@@ -204,7 +192,7 @@ pub fn df_load_csv(
         .with_rechunk(do_rechunk)
         .with_encoding(encoding)
         .with_columns(column_names)
-        .with_dtypes(schema)
+        .with_dtypes(Some(schema_from_dtypes_pairs(dtypes)?))
         .with_null_values(Some(NullValues::AllColumns(null_vals)))
         .with_end_of_line_char(eol_delimiter.unwrap_or(b'\n'));
 

diff --git a/native/explorer/src/lazyframe/io.rs b/native/explorer/src/lazyframe/io.rs
@@ -158,7 +158,7 @@ pub fn lf_from_csv(
     skip_rows: usize,
     delimiter_as_byte: u8,
     do_rechunk: bool,
-    dtypes: Option<Vec<(&str, &str)>>,
+    dtypes: Vec<(&str, &str)>,
     encoding: &str,
     null_vals: Vec<String>,
     parse_dates: bool,
@@ -169,12 +169,6 @@ pub fn lf_from_csv(
         _ => CsvEncoding::Utf8,
     };
 
-    let schema = match dtypes {
-        Some(dtypes) => Some(schema_from_dtypes_pairs(dtypes)?),
-
-        None => None,
-    };
-
     let df = LazyCsvReader::new(filename)
         .with_infer_schema_length(infer_schema_length)
         .has_header(has_header)
@@ -184,7 +178,7 @@ pub fn lf_from_csv(
         .with_skip_rows(skip_rows)
         .with_rechunk(do_rechunk)
         .with_encoding(encoding)
-        .with_dtype_overwrite(schema.as_deref())
+        .with_dtype_overwrite(Some(schema_from_dtypes_pairs(dtypes)?.as_ref()))
         .with_null_values(Some(NullValues::AllColumns(null_vals)))
         .with_end_of_line_char(eol_delimiter.unwrap_or(b'\n'))
         .finish()?;