fix[rust]: csv statistics also sample from end of file (#4554)

polars/polars-io/src/csv/parser.rs

@@ -148,33 +148,49 @@ pub(crate) fn skip_line_ending(input: &[u8], eol_char: u8) -> &[u8] {
 }
 
 /// Get the mean and standard deviation of length of lines in bytes
-pub(crate) fn get_line_stats(bytes: &[u8], n_lines: usize, eol_char: u8) -> Option<(f32, f32)> {
-    let mut n_read = 0;
+pub(crate) fn get_line_stats(
+    bytes: &[u8],
+    n_lines: usize,
+    eol_char: u8,
+    expected_fields: usize,
+    delimiter: u8,
+    quote_char: Option<u8>,
+) -> Option<(f32, f32)> {
     let mut lengths = Vec::with_capacity(n_lines);
-    let file_len = bytes.len();
+
     let mut bytes_trunc;
+    let n_lines_per_iter = n_lines / 2;
 
-    for _ in 0..n_lines {
-        if n_read >= file_len {
-            return None;
-        }
-        bytes_trunc = &bytes[n_read..];
-        match memchr::memchr(eol_char, bytes_trunc) {
-            Some(position) => {
-                n_read += position + 1;
-                lengths.push(position + 1);
-            }
-            None => {
-                return None;
-            }
+    let mut n_read = 0;
+
+    // sample from start and 75% in the file
+    for offset in [0, (bytes.len() as f32 * 0.75) as usize] {
+        bytes_trunc = &bytes[offset..];
+        let pos = next_line_position(
+            bytes_trunc,
+            expected_fields,
+            delimiter,
+            quote_char,
+            eol_char,
+        )?;
+        bytes_trunc = &bytes_trunc[pos + 1..];
+
+        for _ in offset..(offset + n_lines_per_iter) {
+            let pos = next_line_position_naive(bytes_trunc, eol_char)? + 1;
+            n_read += pos;
+            lengths.push(pos);
+            bytes_trunc = &bytes_trunc[pos..];
         }
     }
-    let mean = (n_read as f32) / (n_lines as f32);
+
+    let n_samples = lengths.len();
+
+    let mean = (n_read as f32) / (n_samples as f32);
     let mut std = 0.0;
-    for &len in lengths.iter().take(n_lines) {
+    for &len in lengths.iter() {
         std += (len as f32 - mean).pow(2.0)
     }
-    std = (std / n_lines as f32).pow(0.5);
+    std = (std / n_samples as f32).sqrt();
     Some((mean, std))
 }
 

polars/polars-io/src/csv/read_impl.rs

@@ -363,7 +363,14 @@ impl<'a> CoreReader<'a> {
         let mut total_rows = 128;
 
         // if None, there are less then 128 rows in the file and the statistics don't matter that much
-        if let Some((mean, std)) = get_line_stats(bytes, self.sample_size, self.eol_char) {
+        if let Some((mean, std)) = get_line_stats(
+            bytes,
+            self.sample_size,
+            self.eol_char,
+            self.schema.len(),
+            self.delimiter,
+            self.quote_char,
+        ) {
             if logging {
                 eprintln!("avg line length: {}\nstd. dev. line length: {}", mean, std);
             }

polars/polars-io/src/ndjson_core/ndjson.rs

@@ -178,7 +178,14 @@ impl<'a> CoreJsonReader<'a> {
         let mut bytes = bytes;
         let mut total_rows = 128;
 
-        if let Some((mean, std)) = get_line_stats(bytes, self.sample_size, NEWLINE) {
+        if let Some((mean, std)) = get_line_stats(
+            bytes,
+            self.sample_size,
+            NEWLINE,
+            self.schema.len(),
+            b',',
+            None,
+        ) {
             let line_length_upper_bound = mean + 1.1 * std;
 
             total_rows = (bytes.len() as f32 / (mean - 0.01 * std)) as usize;

py-polars/Makefile

@@ -32,6 +32,7 @@ test-all: venv
 	$(PYTHON_BIN)/maturin develop
 	$(PYTHON) -m pytest tests
 	$(PYTHON) -m pytest tests_parametric
+	$(PYTHON) -m pytest tests_slower
 	$(PYTHON) tests/run_doc_examples.py
 
 test-with-cov: venv

py-polars/tests_slower/test_csv.py

@@ -0,0 +1,11 @@
+import io
+
+import polars as pl
+
+
+def test_csv_statistics_offset() -> None:
+    # this would fail if the statistics sample did not also sample
+    # from the end of the file
+    # the lines at the end have larger rows as the numbers increase
+    csv = "\n".join([str(x) for x in range(5_000)])
+    assert pl.read_csv(io.StringIO(csv), n_rows=5000).height == 4999