rust tests

pola-rs · Jan 18, 2024 · 71ce6d4 · 71ce6d4
1 parent 4c0c708
commit 71ce6d4
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 8 deletions.
diff --git a/crates/polars-core/src/datatypes/any_value.rs b/crates/polars-core/src/datatypes/any_value.rs
@@ -1194,8 +1194,8 @@ mod test {
             ),
             (ArrowDataType::LargeUtf8, DataType::String),
             (ArrowDataType::Utf8, DataType::String),
-            (ArrowDataType::LargeBinary, DataType::BinaryOffset),
-            (ArrowDataType::Binary, DataType::BinaryOffset),
+            (ArrowDataType::LargeBinary, DataType::Binary),
+            (ArrowDataType::Binary, DataType::Binary),
             (
                 ArrowDataType::Time64(ArrowTimeUnit::Nanosecond),
                 DataType::Time,

diff --git a/crates/polars-io/src/csv/read_impl/batched_mmap.rs b/crates/polars-io/src/csv/read_impl/batched_mmap.rs
@@ -129,7 +129,7 @@ impl<'a> CoreReader<'a> {
             eol_char: self.eol_char,
         };
 
-        let projection = self.get_projection();
+        let projection = self.get_projection()?;
 
         // RAII structure that will ensure we maintain a global stringcache
         #[cfg(feature = "dtype-categorical")]

diff --git a/crates/polars-io/src/csv/read_impl/batched_read.rs b/crates/polars-io/src/csv/read_impl/batched_read.rs
@@ -212,7 +212,7 @@ impl<'a> CoreReader<'a> {
             4096,
         );
 
-        let projection = self.get_projection();
+        let projection = self.get_projection()?;
 
         // RAII structure that will ensure we maintain a global stringcache
         #[cfg(feature = "dtype-categorical")]

diff --git a/crates/polars-io/src/csv/read_impl/mod.rs b/crates/polars-io/src/csv/read_impl/mod.rs
@@ -436,16 +436,19 @@ impl<'a> CoreReader<'a> {
             remaining_bytes,
         ))
     }
-    fn get_projection(&mut self) -> Vec<usize> {
+    fn get_projection(&mut self) -> PolarsResult<Vec<usize>> {
         // we also need to sort the projection to have predictable output.
         // the `parse_lines` function expects this.
         self.projection
             .take()
             .map(|mut v| {
                 v.sort_unstable();
-                v
+                if let Some(idx) = v.last() {
+                    polars_ensure!(*idx < self.schema.len(), OutOfBounds: "projection index: {} is out of bounds for csv schema with length: {}", idx, self.schema.len())
+                }
+                Ok(v)
             })
-            .unwrap_or_else(|| (0..self.schema.len()).collect())
+            .unwrap_or_else(|| Ok((0..self.schema.len()).collect()))
     }
 
     fn parse_csv(
@@ -457,7 +460,7 @@ impl<'a> CoreReader<'a> {
         let logging = verbose();
         let (file_chunks, chunk_size, total_rows, starting_point_offset, bytes, remaining_bytes) =
             self.determine_file_chunks_and_statistics(&mut n_threads, bytes, logging)?;
-        let projection = self.get_projection();
+        let projection = self.get_projection()?;
 
         // An empty file with a schema should return an empty DataFrame with that schema
         if bytes.is_empty() {