Allow Index Column of Type Duration in DataFrame.rolling

crates/polars-arrow/src/array/dictionary/mod.rs

@@ -25,7 +25,9 @@ use polars_error::{polars_bail, PolarsResult};
 use super::primitive::PrimitiveArray;
 use super::specification::check_indexes;
 use super::{new_empty_array, new_null_array, Array};
-use crate::array::dictionary::typed_iterator::{DictValue, DictionaryValuesIterTyped};
+use crate::array::dictionary::typed_iterator::{
+    DictValue, DictionaryIterTyped, DictionaryValuesIterTyped,
+};
 
 /// Trait denoting [`NativeType`]s that can be used as keys of a dictionary.
 /// # Safety
@@ -241,30 +243,22 @@ impl<K: DictionaryKey> DictionaryArray<K> {
     ///
     /// # Panics
     ///
-    /// Panics if the keys of this [`DictionaryArray`] have any null types.
-    /// If they do [`DictionaryArray::iter_typed`] should be called
+    /// Panics if the keys of this [`DictionaryArray`] has any nulls.
+    /// If they do [`DictionaryArray::iter_typed`] should be used.
     pub fn values_iter_typed<V: DictValue>(&self) -> PolarsResult<DictionaryValuesIterTyped<K, V>> {
         let keys = &self.keys;
         assert_eq!(keys.null_count(), 0);
         let values = self.values.as_ref();
         let values = V::downcast_values(values)?;
-        Ok(unsafe { DictionaryValuesIterTyped::new(keys, values) })
+        Ok(DictionaryValuesIterTyped::new(keys, values))
     }
 
     /// Returns an iterator over the optional values of  [`Option<V::IterValue>`].
-    ///
-    /// # Panics
-    ///
-    /// This function panics if the `values` array
-    pub fn iter_typed<V: DictValue>(
-        &self,
-    ) -> PolarsResult<ZipValidity<V::IterValue<'_>, DictionaryValuesIterTyped<K, V>, BitmapIter>>
-    {
+    pub fn iter_typed<V: DictValue>(&self) -> PolarsResult<DictionaryIterTyped<K, V>> {
         let keys = &self.keys;
         let values = self.values.as_ref();
         let values = V::downcast_values(values)?;
-        let values_iter = unsafe { DictionaryValuesIterTyped::new(keys, values) };
-        Ok(ZipValidity::new_with_validity(values_iter, self.validity()))
+        Ok(DictionaryIterTyped::new(keys, values))
     }
 
     /// Returns the [`ArrowDataType`] of this [`DictionaryArray`]

crates/polars-arrow/src/array/dictionary/typed_iterator.rs

@@ -1,7 +1,7 @@
 use polars_error::{polars_err, PolarsResult};
 
 use super::DictionaryKey;
-use crate::array::{Array, PrimitiveArray, Utf8Array, Utf8ViewArray};
+use crate::array::{Array, PrimitiveArray, StaticArray, Utf8Array, Utf8ViewArray};
 use crate::trusted_len::TrustedLen;
 use crate::types::Offset;
 
@@ -85,7 +85,8 @@ pub struct DictionaryValuesIterTyped<'a, K: DictionaryKey, V: DictValue> {
 }
 
 impl<'a, K: DictionaryKey, V: DictValue> DictionaryValuesIterTyped<'a, K, V> {
-    pub(super) unsafe fn new(keys: &'a PrimitiveArray<K>, values: &'a V) -> Self {
+    pub(super) fn new(keys: &'a PrimitiveArray<K>, values: &'a V) -> Self {
+        assert_eq!(keys.null_count(), 0);
         Self {
             keys,
             values,
@@ -137,3 +138,68 @@ impl<'a, K: DictionaryKey, V: DictValue> DoubleEndedIterator
         }
     }
 }
+
+pub struct DictionaryIterTyped<'a, K: DictionaryKey, V: DictValue> {
+    keys: &'a PrimitiveArray<K>,
+    values: &'a V,
+    index: usize,
+    end: usize,
+}
+
+impl<'a, K: DictionaryKey, V: DictValue> DictionaryIterTyped<'a, K, V> {
+    pub(super) fn new(keys: &'a PrimitiveArray<K>, values: &'a V) -> Self {
+        Self {
+            keys,
+            values,
+            index: 0,
+            end: keys.len(),
+        }
+    }
+}
+
+impl<'a, K: DictionaryKey, V: DictValue> Iterator for DictionaryIterTyped<'a, K, V> {
+    type Item = Option<V::IterValue<'a>>;
+
+    #[inline]
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.index == self.end {
+            return None;
+        }
+        let old = self.index;
+        self.index += 1;
+        unsafe {
+            if let Some(key) = self.keys.get_unchecked(old) {
+                let idx = key.as_usize();
+                Some(Some(self.values.get_unchecked(idx)))
+            } else {
+                Some(None)
+            }
+        }
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        (self.end - self.index, Some(self.end - self.index))
+    }
+}
+
+unsafe impl<'a, K: DictionaryKey, V: DictValue> TrustedLen for DictionaryIterTyped<'a, K, V> {}
+
+impl<'a, K: DictionaryKey, V: DictValue> DoubleEndedIterator for DictionaryIterTyped<'a, K, V> {
+    #[inline]
+    fn next_back(&mut self) -> Option<Self::Item> {
+        if self.index == self.end {
+            None
+        } else {
+            self.end -= 1;
+            unsafe {
+                if let Some(key) = self.keys.get_unchecked(self.end) {
+                    let idx = key.as_usize();
+                    Some(Some(self.values.get_unchecked(idx)))
+                } else {
+                    Some(None)
+                }
+            }
+        }
+    }
+}

crates/polars-core/src/chunked_array/logical/categorical/mod.rs

@@ -279,6 +279,10 @@ impl CategoricalChunked {
         self
     }
 
+    pub fn _with_fast_unique(self, toggle: bool) -> Self {
+        self.with_fast_unique(toggle)
+    }
+
     /// Get a reference to the mapping of categorical types to the string values.
     pub fn get_rev_map(&self) -> &Arc<RevMapping> {
         if let DataType::Categorical(Some(rev_map), _) | DataType::Enum(Some(rev_map), _) =

crates/polars-core/src/utils/supertype.rs

@@ -264,28 +264,43 @@ pub fn get_supertype(l: &DataType, r: &DataType) -> Option<DataType> {
             },
             (dt, Unknown(kind)) => {
                 match kind {
+                    // numeric vs float|str -> always float|str
                     UnknownKind::Float | UnknownKind::Int(_) if dt.is_float() | dt.is_string() => Some(dt.clone()),
-                    UnknownKind::Float if dt.is_numeric() => Some(Unknown(UnknownKind::Float)),
+                    UnknownKind::Float if dt.is_integer() => Some(Unknown(UnknownKind::Float)),
+                    // Materialize float
+                    UnknownKind::Float if dt.is_float() => Some(dt.clone()),
+                    // Materialize str
                     UnknownKind::Str if dt.is_string() | dt.is_enum() => Some(dt.clone()),
+                    // Materialize str
                     #[cfg(feature = "dtype-categorical")]
                     UnknownKind::Str if dt.is_categorical()  => {
                         let Categorical(_, ord) = dt else { unreachable!()};
                         Some(Categorical(None, *ord))
                     },
+                    // Keep unknown
                     dynam if dt.is_null() => Some(Unknown(*dynam)),
+                    // Find integers sizes
                     UnknownKind::Int(v) if dt.is_numeric() => {
-                        let smallest_fitting_dtype = if dt.is_unsigned_integer() && v.is_positive() {
-                            materialize_dyn_int_pos(*v).dtype()
-                        } else {
-                            materialize_smallest_dyn_int(*v).dtype()
-                        };
-                        match dt {
-                            UInt64 if smallest_fitting_dtype.is_signed_integer() => {
-                                // Ensure we don't cast to float when dealing with dynamic literals
-                                Some(Int64)
-                            },
-                            _ => {
-                                get_supertype(dt, &smallest_fitting_dtype)
+                        // Both dyn int
+                        if let Unknown(UnknownKind::Int(v_other)) = dt {
+                            // Take the maximum value to ensure we bubble up the required minimal size.
+                            Some(Unknown(UnknownKind::Int(std::cmp::max(*v, *v_other))))
+                        }
+                        // dyn int vs number
+                        else {
+                            let smallest_fitting_dtype = if dt.is_unsigned_integer() && v.is_positive() {
+                                materialize_dyn_int_pos(*v).dtype()
+                            } else {
+                                materialize_smallest_dyn_int(*v).dtype()
+                            };
+                            match dt {
+                                UInt64 if smallest_fitting_dtype.is_signed_integer() => {
+                                    // Ensure we don't cast to float when dealing with dynamic literals
+                                    Some(Int64)
+                                },
+                                _ => {
+                                    get_supertype(dt, &smallest_fitting_dtype)
+                                }
                             }
                         }
                     }

crates/polars-io/src/csv/read/read_impl/batched_read.rs

@@ -7,6 +7,7 @@ use polars_core::frame::DataFrame;
 use polars_core::schema::SchemaRef;
 use polars_core::POOL;
 use polars_error::PolarsResult;
+use polars_utils::sync::SyncPtr;
 use polars_utils::IdxSize;
 use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
 
@@ -54,6 +55,8 @@ pub(crate) fn get_offsets(
     }
 }
 
+/// Reads bytes from `file` to `buf` and returns pointers into `buf` that can be parsed.
+/// TODO! this can be implemented without copying by pointing in the memmapped file.
 struct ChunkReader<'a> {
     file: &'a File,
     buf: Vec<u8>,
@@ -109,18 +112,23 @@ impl<'a> ChunkReader<'a> {
         self.buf_end = 0;
     }
 
-    fn return_slice(&self, start: usize, end: usize) -> (usize, usize) {
+    fn return_slice(&self, start: usize, end: usize) -> (SyncPtr<u8>, usize) {
         let slice = &self.buf[start..end];
         let len = slice.len();
-        (slice.as_ptr() as usize, len)
+        (slice.as_ptr().into(), len)
     }
 
-    fn get_buf(&self) -> (usize, usize) {
+    fn get_buf_remaining(&self) -> (SyncPtr<u8>, usize) {
         let slice = &self.buf[self.buf_end..];
         let len = slice.len();
-        (slice.as_ptr() as usize, len)
+        (slice.as_ptr().into(), len)
     }
 
+    // Get next `n` offset positions. Where `n` is number of chunks.
+
+    // This returns pointers into slices into `buf`
+    // we must process the slices before the next call
+    // as that will overwrite the slices
     fn read(&mut self, n: usize) -> bool {
         self.reslice();
 
@@ -267,7 +275,7 @@ pub struct BatchedCsvReaderRead<'a> {
     chunk_size: usize,
     finished: bool,
     file_chunk_reader: ChunkReader<'a>,
-    file_chunks: Vec<(usize, usize)>,
+    file_chunks: Vec<(SyncPtr<u8>, usize)>,
     projection: Vec<usize>,
     starting_point_offset: Option<usize>,
     row_index: Option<RowIndex>,
@@ -292,6 +300,7 @@ pub struct BatchedCsvReaderRead<'a> {
 }
 //
 impl<'a> BatchedCsvReaderRead<'a> {
+    /// `n` number of batches.
     pub fn next_batches(&mut self, n: usize) -> PolarsResult<Option<Vec<DataFrame>>> {
         if n == 0 || self.finished {
             return Ok(None);
@@ -320,7 +329,8 @@ impl<'a> BatchedCsvReaderRead<'a> {
         // ensure we process the final slice as well.
         if self.file_chunk_reader.finished && self.file_chunks.len() < n {
             // get the final slice
-            self.file_chunks.push(self.file_chunk_reader.get_buf());
+            self.file_chunks
+                .push(self.file_chunk_reader.get_buf_remaining());
             self.finished = true
         }
 
@@ -333,7 +343,7 @@ impl<'a> BatchedCsvReaderRead<'a> {
             self.file_chunks
                 .par_iter()
                 .map(|(ptr, len)| {
-                    let chunk = unsafe { std::slice::from_raw_parts(*ptr as *const u8, *len) };
+                    let chunk = unsafe { std::slice::from_raw_parts(ptr.get(), *len) };
                     let stop_at_n_bytes = chunk.len();
                     let mut df = read_chunk(
                         chunk,