Merge remote-tracking branch 'upstream/main' into round-durations

pola-rs · Apr 18, 2024 · 3c3147b · 3c3147b
2 parents f8edaae + 96e1f01
commit 3c3147b
Show file tree

Hide file tree

Showing 16 changed files with 302 additions and 208 deletions.
diff --git a/crates/polars-core/src/chunked_array/cast.rs b/crates/polars-core/src/chunked_array/cast.rs
@@ -136,9 +136,6 @@ where
                         polars_bail!(OutOfBounds: "index {} is bigger than the number of categories {}",m,categories.len());
                     }
                 }
-                // SAFETY:
-                // we are guarded by the type system
-                let ca = unsafe { &*(self as *const ChunkedArray<T> as *const UInt32Chunked) };
                 // SAFETY: indices are in bound
                 unsafe {
                     Ok(CategoricalChunked::from_cats_and_rev_map_unchecked(

diff --git a/crates/polars-lazy/src/scan/csv.rs b/crates/polars-lazy/src/scan/csv.rs
@@ -10,7 +10,7 @@ use crate::prelude::*;
 
 #[derive(Clone)]
 #[cfg(feature = "csv")]
-pub struct LazyCsvReader<'a> {
+pub struct LazyCsvReader {
     path: PathBuf,
     paths: Arc<[PathBuf]>,
     separator: u8,
@@ -20,7 +20,7 @@ pub struct LazyCsvReader<'a> {
     n_rows: Option<usize>,
     cache: bool,
     schema: Option<SchemaRef>,
-    schema_overwrite: Option<&'a Schema>,
+    schema_overwrite: Option<SchemaRef>,
     low_memory: bool,
     comment_prefix: Option<CommentPrefix>,
     quote_char: Option<u8>,
@@ -39,7 +39,7 @@ pub struct LazyCsvReader<'a> {
 }
 
 #[cfg(feature = "csv")]
-impl<'a> LazyCsvReader<'a> {
+impl LazyCsvReader {
     pub fn new_paths(paths: Arc<[PathBuf]>) -> Self {
         Self::new("").with_paths(paths)
     }
@@ -129,7 +129,7 @@ impl<'a> LazyCsvReader<'a> {
     /// Overwrite the schema with the dtypes in this given Schema. The given schema may be a subset
     /// of the total schema.
     #[must_use]
-    pub fn with_dtype_overwrite(mut self, schema: Option<&'a Schema>) -> Self {
+    pub fn with_dtype_overwrite(mut self, schema: Option<SchemaRef>) -> Self {
         self.schema_overwrite = schema;
         self
     }
@@ -270,7 +270,7 @@ impl<'a> LazyCsvReader<'a> {
         let mut schema = f(schema)?;
 
         // the dtypes set may be for the new names, so update again
-        if let Some(overwrite_schema) = self.schema_overwrite {
+        if let Some(overwrite_schema) = &self.schema_overwrite {
             for (name, dtype) in overwrite_schema.iter() {
                 schema.with_column(name.clone(), dtype.clone());
             }
@@ -280,7 +280,7 @@ impl<'a> LazyCsvReader<'a> {
     }
 }
 
-impl LazyFileListReader for LazyCsvReader<'_> {
+impl LazyFileListReader for LazyCsvReader {
     fn finish_no_glob(self) -> PolarsResult<LazyFrame> {
         let mut lf: LazyFrame = DslBuilder::scan_csv(
             self.path,

diff --git a/crates/polars-lazy/src/tests/io.rs b/crates/polars-lazy/src/tests/io.rs
@@ -672,10 +672,10 @@ fn scan_small_dtypes() -> PolarsResult<()> {
     for dt in small_dt {
         let df = LazyCsvReader::new(FOODS_CSV)
             .has_header(true)
-            .with_dtype_overwrite(Some(&Schema::from_iter([Field::new(
+            .with_dtype_overwrite(Some(Arc::new(Schema::from_iter([Field::new(
                 "sugars_g",
                 dt.clone(),
-            )])))
+            )]))))
             .finish()?
             .select(&[col("sugars_g")])
             .collect()?;

diff --git a/crates/polars-plan/src/dot.rs b/crates/polars-plan/src/dot.rs
@@ -199,7 +199,7 @@ impl DslPlan {
                 "PYTHON",
                 &[],
                 options.with_columns.as_ref().map(|s| s.as_slice()),
-                options.schema.len(),
+                Some(options.schema.len()),
                 &options.predicate,
                 branch,
                 id,
@@ -342,7 +342,7 @@ impl DslPlan {
                     name,
                     paths.as_ref(),
                     options.with_columns.as_ref().map(|cols| cols.as_slice()),
-                    file_info.schema.len(),
+                    file_info.as_ref().map(|fi| fi.schema.len()),
                     predicate,
                     branch,
                     id,
@@ -418,7 +418,7 @@ impl DslPlan {
         name: &str,
         path: &[PathBuf],
         with_columns: Option<&[String]>,
-        total_columns: usize,
+        total_columns: Option<usize>,
         predicate: &Option<P>,
         branch: usize,
         id: usize,
@@ -440,6 +440,9 @@ impl DslPlan {
         };
 
         let pred = fmt_predicate(predicate.as_ref());
+        let total_columns = total_columns
+            .map(|v| format!("{v}"))
+            .unwrap_or_else(|| "?".to_string());
         let fmt = format!(
             "{name} SCAN {};\nπ {}/{};\nσ {}",
             path_fmt, n_columns_fmt, total_columns, pred,