don't copy the sorted flag on many operations (#3795)

pola-rs · Jun 24, 2022 · c578bfb · c578bfb
1 parent b59b6d7
commit c578bfb
Show file tree

Hide file tree

Showing 25 changed files with 149 additions and 91 deletions.
diff --git a/polars/polars-core/src/chunked_array/arithmetic.rs b/polars/polars-core/src/chunked_array/arithmetic.rs
@@ -70,7 +70,7 @@ where
                 .zip(rhs.downcast_iter())
                 .map(|(lhs, rhs)| Box::new(kernel(lhs, rhs)) as ArrayRef)
                 .collect();
-            lhs.copy_with_chunks(chunks)
+            lhs.copy_with_chunks(chunks, false)
         }
         // broadcast right path
         (_, 1) => {

diff --git a/polars/polars-core/src/chunked_array/cast.rs b/polars/polars-core/src/chunked_array/cast.rs
@@ -64,18 +64,13 @@ where
             }
             _ => cast_impl_inner(self.name(), &self.chunks, data_type, checked).map(|mut s| {
                 // maintain sorted if data types remain signed
-                if self.is_sorted()
-                    || self.is_sorted_reverse() && (s.null_count() == self.null_count())
+                // this may still fail with overflow?
+                if ((self.dtype().is_signed() && data_type.is_signed())
+                    || (self.dtype().is_unsigned() && data_type.is_unsigned()))
+                    && (s.null_count() == self.null_count())
                 {
-                    // this may still fail with overflow?
-                    //
-                    if (self.dtype().is_signed() && data_type.is_signed())
-                        || (self.dtype().is_unsigned() && data_type.is_unsigned())
-                    {
-                        let reverse = self.is_sorted_reverse();
-                        let inner = s._get_inner_mut();
-                        inner._set_sorted(reverse)
-                    }
+                    let is_sorted = self.is_sorted2();
+                    s.set_sorted(is_sorted)
                 }
                 s
             }),

diff --git a/polars/polars-core/src/chunked_array/mod.rs b/polars/polars-core/src/chunked_array/mod.rs
@@ -177,6 +177,16 @@ impl<T> ChunkedArray<T> {
         }
     }
 
+    pub fn is_sorted2(&self) -> IsSorted {
+        if self.is_sorted() {
+            IsSorted::Ascending
+        } else if self.is_sorted_reverse() {
+            IsSorted::Descending
+        } else {
+            IsSorted::Not
+        }
+    }
+
     pub fn set_sorted2(&mut self, sorted: IsSorted) {
         match sorted {
             IsSorted::Not => {
@@ -325,14 +335,18 @@ impl<T> ChunkedArray<T> {
     }
 
     /// Create a new ChunkedArray from self, where the chunks are replaced.
-    fn copy_with_chunks(&self, chunks: Vec<ArrayRef>) -> Self {
-        ChunkedArray {
+    fn copy_with_chunks(&self, chunks: Vec<ArrayRef>, keep_sorted: bool) -> Self {
+        let mut out = ChunkedArray {
             field: self.field.clone(),
             chunks,
             phantom: PhantomData,
             categorical_map: self.categorical_map.clone(),
             bit_settings: self.bit_settings,
+        };
+        if !keep_sorted {
+            out.set_sorted2(IsSorted::Not);
         }
+        out
     }
 
     /// Get a mask of the null values.

diff --git a/polars/polars-core/src/chunked_array/ops/chunkops.rs b/polars/polars-core/src/chunked_array/ops/chunkops.rs
@@ -70,7 +70,7 @@ where
     }
     #[inline]
     fn slice(&self, offset: i64, length: usize) -> Self {
-        self.copy_with_chunks(slice(&self.chunks, offset, length, self.len()))
+        self.copy_with_chunks(slice(&self.chunks, offset, length, self.len()), true)
     }
 }
 
@@ -92,7 +92,7 @@ impl ChunkOps for BooleanChunked {
     }
     #[inline]
     fn slice(&self, offset: i64, length: usize) -> Self {
-        self.copy_with_chunks(slice(&self.chunks, offset, length, self.len()))
+        self.copy_with_chunks(slice(&self.chunks, offset, length, self.len()), true)
     }
 }
 
@@ -109,12 +109,12 @@ impl ChunkOps for Utf8Chunked {
                     .as_slice(),
             )
             .unwrap()];
-            self.copy_with_chunks(chunks)
+            self.copy_with_chunks(chunks, true)
         }
     }
     #[inline]
     fn slice(&self, offset: i64, length: usize) -> Self {
-        self.copy_with_chunks(slice(&self.chunks, offset, length, self.len()))
+        self.copy_with_chunks(slice(&self.chunks, offset, length, self.len()), true)
     }
 }
 
@@ -140,7 +140,7 @@ impl ChunkOps for ListChunked {
     }
     #[inline]
     fn slice(&self, offset: i64, length: usize) -> Self {
-        self.copy_with_chunks(slice(&self.chunks, offset, length, self.len()))
+        self.copy_with_chunks(slice(&self.chunks, offset, length, self.len()), true)
     }
 }
 
@@ -183,7 +183,7 @@ where
     }
     #[inline]
     fn slice(&self, offset: i64, length: usize) -> Self {
-        self.copy_with_chunks(slice(&self.chunks, offset, length, self.len()))
+        self.copy_with_chunks(slice(&self.chunks, offset, length, self.len()), true)
     }
 }
 

diff --git a/polars/polars-core/src/chunked_array/ops/filter.rs b/polars/polars-core/src/chunked_array/ops/filter.rs
@@ -45,7 +45,7 @@ where
             .zip(filter.downcast_iter())
             .map(|(left, mask)| filter_fn(left, mask).unwrap())
             .collect::<Vec<_>>();
-        Ok(self.copy_with_chunks(chunks))
+        Ok(self.copy_with_chunks(chunks, true))
     }
 }
 
@@ -66,7 +66,7 @@ impl ChunkFilter<BooleanType> for BooleanChunked {
             .zip(filter.downcast_iter())
             .map(|(left, mask)| filter_fn(left, mask).unwrap())
             .collect::<Vec<_>>();
-        Ok(self.copy_with_chunks(chunks))
+        Ok(self.copy_with_chunks(chunks, true))
     }
 }
 
@@ -88,7 +88,7 @@ impl ChunkFilter<Utf8Type> for Utf8Chunked {
             .map(|(left, mask)| filter_fn(left, mask).unwrap())
             .collect::<Vec<_>>();
 
-        Ok(self.copy_with_chunks(chunks))
+        Ok(self.copy_with_chunks(chunks, true))
     }
 }
 

diff --git a/polars/polars-core/src/chunked_array/ops/take/mod.rs b/polars/polars-core/src/chunked_array/ops/take/mod.rs
@@ -101,7 +101,7 @@ where
                         }
                     }
                 };
-                self.copy_with_chunks(vec![array])
+                self.copy_with_chunks(vec![array], false)
             }
             TakeIdx::Iter(iter) => {
                 if self.is_empty() {
@@ -122,7 +122,7 @@ where
                         return ca;
                     }
                 };
-                self.copy_with_chunks(vec![array])
+                self.copy_with_chunks(vec![array], false)
             }
             TakeIdx::IterNulls(iter) => {
                 if self.is_empty() {
@@ -143,7 +143,7 @@ where
                         return ca;
                     }
                 };
-                self.copy_with_chunks(vec![array])
+                self.copy_with_chunks(vec![array], false)
             }
         }
     }
@@ -192,7 +192,7 @@ impl ChunkTake for BooleanChunked {
                         }
                     }
                 };
-                self.copy_with_chunks(vec![array])
+                self.copy_with_chunks(vec![array], false)
             }
             TakeIdx::Iter(iter) => {
                 if self.is_empty() {
@@ -209,7 +209,7 @@ impl ChunkTake for BooleanChunked {
                         return ca;
                     }
                 };
-                self.copy_with_chunks(vec![array])
+                self.copy_with_chunks(vec![array], false)
             }
             TakeIdx::IterNulls(iter) => {
                 if self.is_empty() {
@@ -229,7 +229,7 @@ impl ChunkTake for BooleanChunked {
                         return ca;
                     }
                 };
-                self.copy_with_chunks(vec![array])
+                self.copy_with_chunks(vec![array], false)
             }
         }
     }
@@ -278,7 +278,7 @@ impl ChunkTake for Utf8Chunked {
                         }
                     }
                 };
-                self.copy_with_chunks(vec![array])
+                self.copy_with_chunks(vec![array], false)
             }
             TakeIdx::Iter(iter) => {
                 let array = match (self.has_validity(), self.chunks.len()) {
@@ -292,7 +292,7 @@ impl ChunkTake for Utf8Chunked {
                         return ca;
                     }
                 };
-                self.copy_with_chunks(vec![array])
+                self.copy_with_chunks(vec![array], false)
             }
             TakeIdx::IterNulls(iter) => {
                 let array = match (self.has_validity(), self.chunks.len()) {
@@ -309,7 +309,7 @@ impl ChunkTake for Utf8Chunked {
                         return ca;
                     }
                 };
-                self.copy_with_chunks(vec![array])
+                self.copy_with_chunks(vec![array], false)
             }
         }
     }
@@ -369,7 +369,7 @@ impl ChunkTake for ListChunked {
                         }
                     }
                 };
-                ca_self.copy_with_chunks(vec![array])
+                ca_self.copy_with_chunks(vec![array], false)
             }
             // todo! fast path for single chunk
             TakeIdx::Iter(iter) => {

diff --git a/polars/polars-core/src/frame/groupby/mod.rs b/polars/polars-core/src/frame/groupby/mod.rs
@@ -332,7 +332,11 @@ impl<'df> GroupBy<'df> {
                             let mut iter = groups.iter().map(|(first, _idx)| first as usize);
                             // Safety:
                             // groups are always in bounds
-                            unsafe { s.take_iter_unchecked(&mut iter) }
+                            let mut out = unsafe { s.take_iter_unchecked(&mut iter) };
+                            if groups.sorted {
+                                out.set_sorted(s.is_sorted());
+                            };
+                            out
                         }
                         GroupsProxy::Slice { groups, rolling } => {
                             if *rolling && !groups.is_empty() {
@@ -348,7 +352,10 @@ impl<'df> GroupBy<'df> {
                             let mut iter = groups.iter().map(|&[first, _len]| first as usize);
                             // Safety:
                             // groups are always in bounds
-                            unsafe { s.take_iter_unchecked(&mut iter) }
+                            let mut out = unsafe { s.take_iter_unchecked(&mut iter) };
+                            // sliced groups are always in order of discovery
+                            out.set_sorted(s.is_sorted());
+                            out
                         }
                     }
                 })

diff --git a/polars/polars-core/src/frame/groupby/proxy.rs b/polars/polars-core/src/frame/groupby/proxy.rs
@@ -192,7 +192,12 @@ pub type GroupsSlice = Vec<[IdxSize; 2]>;
 #[derive(Debug, Clone, PartialEq)]
 pub enum GroupsProxy {
     Idx(GroupsIdx),
-    Slice { groups: GroupsSlice, rolling: bool },
+    Slice {
+        // the groups slices
+        groups: GroupsSlice,
+        // indicates if we do a rolling groupby
+        rolling: bool,
+    },
 }
 
 impl Default for GroupsProxy {

diff --git a/polars/polars-core/src/frame/mod.rs b/polars/polars-core/src/frame/mod.rs
@@ -1766,7 +1766,11 @@ impl DataFrame {
         // not present in the dataframe
         let _ = df.apply(&first_by_column, |s| {
             let mut s = s.clone();
-            s.set_sorted(first_reverse);
+            if first_reverse {
+                s.set_sorted(IsSorted::Descending)
+            } else {
+                s.set_sorted(IsSorted::Ascending)
+            }
             s
         });
         Ok(df)

diff --git a/polars/polars-core/src/series/implementations/boolean.rs b/polars/polars-core/src/series/implementations/boolean.rs
@@ -38,8 +38,8 @@ impl private::PrivateSeries for SeriesWrap<BooleanChunked> {
         self.0.explode_by_offsets(offsets)
     }
 
-    fn _set_sorted(&mut self, reverse: bool) {
-        self.0.set_sorted(reverse)
+    fn _set_sorted(&mut self, is_sorted: IsSorted) {
+        self.0.set_sorted2(is_sorted)
     }
 
     unsafe fn equal_element(&self, idx_self: usize, idx_other: usize, other: &Series) -> bool {

diff --git a/polars/polars-core/src/series/implementations/categorical.rs b/polars/polars-core/src/series/implementations/categorical.rs
@@ -72,8 +72,8 @@ impl private::PrivateSeries for SeriesWrap<CategoricalChunked> {
         .into_series()
     }
 
-    fn _set_sorted(&mut self, reverse: bool) {
-        self.0.logical_mut().set_sorted(reverse)
+    fn _set_sorted(&mut self, is_sorted: IsSorted) {
+        self.0.logical_mut().set_sorted2(is_sorted)
     }
 
     unsafe fn equal_element(&self, idx_self: usize, idx_other: usize, other: &Series) -> bool {

diff --git a/polars/polars-core/src/series/implementations/dates_time.rs b/polars/polars-core/src/series/implementations/dates_time.rs
@@ -57,8 +57,8 @@ macro_rules! impl_dyn_series {
                 self.0.cummin(reverse).$into_logical().into_series()
             }
 
-            fn _set_sorted(&mut self, reverse: bool) {
-                self.0.deref_mut().set_sorted(reverse)
+            fn _set_sorted(&mut self, is_sorted: IsSorted) {
+                self.0.deref_mut().set_sorted2(is_sorted)
             }
 
             #[cfg(feature = "zip_with")]
@@ -297,7 +297,7 @@ macro_rules! impl_dyn_series {
                 let mut out = ChunkTake::take_unchecked(self.0.deref(), idx.into());
 
                 if self.0.is_sorted() && (idx.is_sorted() || idx.is_sorted_reverse()) {
-                    out.set_sorted(idx.is_sorted_reverse())
+                    out.set_sorted2(idx.is_sorted2())
                 }
 
                 Ok(out.$into_logical().into_series())

diff --git a/polars/polars-core/src/series/implementations/datetime.rs b/polars/polars-core/src/series/implementations/datetime.rs
@@ -57,8 +57,8 @@ impl private::PrivateSeries for SeriesWrap<DatetimeChunked> {
             .into_series()
     }
 
-    fn _set_sorted(&mut self, reverse: bool) {
-        self.0.deref_mut().set_sorted(reverse)
+    fn _set_sorted(&mut self, is_sorted: IsSorted) {
+        self.0.deref_mut().set_sorted2(is_sorted)
     }
 
     #[cfg(feature = "zip_with")]
@@ -332,7 +332,7 @@ impl SeriesTrait for SeriesWrap<DatetimeChunked> {
         let mut out = ChunkTake::take_unchecked(self.0.deref(), idx.into());
 
         if self.0.is_sorted() && (idx.is_sorted() || idx.is_sorted_reverse()) {
-            out.set_sorted(idx.is_sorted_reverse())
+            out.set_sorted2(idx.is_sorted2())
         }
 
         Ok(out

diff --git a/polars/polars-core/src/series/implementations/duration.rs b/polars/polars-core/src/series/implementations/duration.rs
@@ -57,8 +57,8 @@ impl private::PrivateSeries for SeriesWrap<DurationChunked> {
             .into_series()
     }
 
-    fn _set_sorted(&mut self, reverse: bool) {
-        self.0.deref_mut().set_sorted(reverse)
+    fn _set_sorted(&mut self, is_sorted: IsSorted) {
+        self.0.deref_mut().set_sorted2(is_sorted)
     }
 
     unsafe fn equal_element(&self, idx_self: usize, idx_other: usize, other: &Series) -> bool {
@@ -359,7 +359,7 @@ impl SeriesTrait for SeriesWrap<DurationChunked> {
         let mut out = ChunkTake::take_unchecked(self.0.deref(), idx.into());
 
         if self.0.is_sorted() && (idx.is_sorted() || idx.is_sorted_reverse()) {
-            out.set_sorted(idx.is_sorted_reverse())
+            out.set_sorted2(idx.is_sorted2())
         }
 
         Ok(out.into_duration(self.0.time_unit()).into_series())

diff --git a/polars/polars-core/src/series/implementations/floats.rs b/polars/polars-core/src/series/implementations/floats.rs
@@ -52,8 +52,8 @@ macro_rules! impl_dyn_series {
                 self.0.cummin(reverse).into_series()
             }
 
-            fn _set_sorted(&mut self, reverse: bool) {
-                self.0.set_sorted(reverse)
+            fn _set_sorted(&mut self, is_sorted: IsSorted) {
+                self.0.set_sorted2(is_sorted)
             }
 
             unsafe fn equal_element(
@@ -279,7 +279,7 @@ macro_rules! impl_dyn_series {
                 let mut out = ChunkTake::take_unchecked(&self.0, (&*idx).into());
 
                 if self.0.is_sorted() && (idx.is_sorted() || idx.is_sorted_reverse()) {
-                    out.set_sorted(idx.is_sorted_reverse())
+                    out.set_sorted2(idx.is_sorted2())
                 }
 
                 Ok(out.into_series())