csv: improve float inference regex

pola-rs · Jan 8, 2022 · 07b350e · 07b350e
1 parent 7788875
commit 07b350e
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git a/polars/polars-io/src/csv.rs b/polars/polars-io/src/csv.rs
@@ -1396,4 +1396,13 @@ A3,\"B4_\"\"with_embedded_double_quotes\"\"\",C4,4";
         assert_eq!(df.dtypes(), expected);
         Ok(())
     }
+
+    #[test]
+    fn test_comma_separated_field_in_tsv() -> Result<()> {
+        let csv = "first\tsecond\n1\t2.3,2.4\n3\t4.5,4.6\n";
+        let file = Cursor::new(csv);
+        let df = CsvReader::new(file).with_delimiter(b'\t').finish()?;
+        assert_eq!(df.dtypes(), &[DataType::Int64, DataType::Utf8]);
+        Ok(())
+    }
 }
diff --git a/polars/polars-io/src/csv_core/utils.rs b/polars/polars-io/src/csv_core/utils.rs
@@ -73,7 +73,8 @@ pub fn get_reader_bytes<R: Read + MmapBytesReader>(reader: &mut R) -> Result<Rea
 
 lazy_static! {
     static ref FLOAT_RE: Regex =
-        Regex::new(r"^\s*-?((\d*\.\d+)[eE]?[-\+]?\d*)|inf|NaN|\d+[eE][-+]\d+$").unwrap();
+        Regex::new(r"^(\s*-?((\d*\.\d+)[eE]?[-\+]?\d*)|[-+]?inf|[-+]?NaN|\d+[eE][-+]\d+)$")
+            .unwrap();
     static ref INTEGER_RE: Regex = Regex::new(r"^\s*-?(\d+)$").unwrap();
     static ref BOOLEAN_RE: Regex = RegexBuilder::new(r"^\s*(true)$|^(false)$")
         .case_insensitive(true)