fast path for explicitly sorted dataframes and window expressions (#2613

)
pola-rs · Feb 11, 2022 · d27546e · d27546e
1 parent 5acad95
commit d27546e
Show file tree

Hide file tree

Showing 17 changed files with 272 additions and 35 deletions.
diff --git a/polars/polars-core/src/frame/groupby/hashing.rs b/polars/polars-core/src/frame/groupby/hashing.rs
@@ -17,7 +17,9 @@ fn finish_group_order(out: Vec<Vec<IdxItem>>, sorted: bool) -> GroupsProxy {
     if sorted {
         let mut out = out.into_iter().flatten().collect::<Vec<_>>();
         out.sort_unstable_by_key(|g| g.0);
-        GroupsProxy::Idx(GroupsIdx::from_iter(out.into_iter()))
+        let mut idx = GroupsIdx::from_iter(out.into_iter());
+        idx.sorted = true;
+        GroupsProxy::Idx(idx)
     } else {
         GroupsProxy::Idx(GroupsIdx::from(out))
     }
@@ -54,7 +56,9 @@ where
             .map(|(_k, v)| v)
             .collect_trusted::<Vec<_>>();
         groups.sort_unstable_by_key(|g| g.0);
-        GroupsProxy::Idx(groups.into_iter().collect())
+        let mut idx: GroupsIdx = groups.into_iter().collect();
+        idx.sorted = true;
+        GroupsProxy::Idx(idx)
     } else {
         GroupsProxy::Idx(hash_tbl.into_iter().map(|(_k, v)| v).collect())
     }

diff --git a/polars/polars-core/src/frame/groupby/proxy.rs b/polars/polars-core/src/frame/groupby/proxy.rs
@@ -9,6 +9,7 @@ use rayon::prelude::*;
 /// this make sorting fast.
 #[derive(Debug, Clone, PartialEq, Eq, Default)]
 pub struct GroupsIdx {
+    pub(crate) sorted: bool,
     first: Vec<u32>,
     all: Vec<Vec<u32>>,
 }
@@ -69,7 +70,12 @@ impl GroupsIdx {
         let (first, all) = POOL.install(|| rayon::join(take_first, take_all));
         self.first = first;
         self.all = all;
+        self.sorted = true
     }
+    pub fn is_sorted(&self) -> bool {
+        self.sorted
+    }
+
     pub fn iter(
         &self,
     ) -> std::iter::Zip<std::iter::Copied<std::slice::Iter<u32>>, std::slice::Iter<Vec<u32>>> {
@@ -98,7 +104,11 @@ impl GroupsIdx {
 impl FromIterator<IdxItem> for GroupsIdx {
     fn from_iter<T: IntoIterator<Item = IdxItem>>(iter: T) -> Self {
         let (first, all) = iter.into_iter().unzip();
-        GroupsIdx { first, all }
+        GroupsIdx {
+            sorted: false,
+            first,
+            all,
+        }
     }
 }
 
@@ -131,7 +141,11 @@ impl FromParallelIterator<IdxItem> for GroupsIdx {
         I: IntoParallelIterator<Item = IdxItem>,
     {
         let (first, all) = par_iter.into_par_iter().unzip();
-        GroupsIdx { first, all }
+        GroupsIdx {
+            sorted: false,
+            first,
+            all,
+        }
     }
 }
 

diff --git a/polars/polars-core/src/frame/mod.rs b/polars/polars-core/src/frame/mod.rs
@@ -3043,7 +3043,7 @@ mod test {
 
     #[test]
     #[cfg_attr(miri, ignore)]
-    fn drop_duplicates() {
+    fn distinct() {
         let df = df! {
             "flt" => [1., 1., 2., 2., 3., 3.],
             "int" => [1, 1, 2, 2, 3, 3, ],
@@ -3052,7 +3052,7 @@ mod test {
         .unwrap();
         dbg!(&df);
         let df = df
-            .drop_duplicates(true, None)
+            .distinct_stable(None, DistinctKeepStrategy::First)
             .unwrap()
             .sort(["flt"], false)
             .unwrap();

diff --git a/polars/polars-core/src/series/implementations/boolean.rs b/polars/polars-core/src/series/implementations/boolean.rs
@@ -1,6 +1,7 @@
 use super::private;
 use super::IntoSeries;
 use super::SeriesTrait;
+use super::*;
 use crate::chunked_array::comparison::*;
 use crate::chunked_array::{
     ops::{
@@ -107,6 +108,16 @@ impl private::PrivateSeries for SeriesWrap<BooleanChunked> {
 }
 
 impl SeriesTrait for SeriesWrap<BooleanChunked> {
+    fn is_sorted(&self) -> IsSorted {
+        if self.0.is_sorted() {
+            IsSorted::Ascending
+        } else if self.0.is_sorted_reverse() {
+            IsSorted::Descending
+        } else {
+            IsSorted::Not
+        }
+    }
+
     #[cfg(feature = "interpolate")]
     fn interpolate(&self) -> Series {
         self.0.clone().into_series()

diff --git a/polars/polars-core/src/series/implementations/categorical.rs b/polars/polars-core/src/series/implementations/categorical.rs
@@ -1,6 +1,7 @@
 use super::private;
 use super::IntoSeries;
 use super::SeriesTrait;
+use super::*;
 use crate::chunked_array::comparison::*;
 use crate::chunked_array::{
     ops::{
@@ -111,6 +112,16 @@ impl private::PrivateSeries for SeriesWrap<CategoricalChunked> {
 }
 
 impl SeriesTrait for SeriesWrap<CategoricalChunked> {
+    fn is_sorted(&self) -> IsSorted {
+        if self.0.is_sorted() {
+            IsSorted::Ascending
+        } else if self.0.is_sorted_reverse() {
+            IsSorted::Descending
+        } else {
+            IsSorted::Not
+        }
+    }
+
     #[cfg(feature = "interpolate")]
     fn interpolate(&self) -> Series {
         self.0.interpolate().into_series()

diff --git a/polars/polars-core/src/series/implementations/dates_time.rs b/polars/polars-core/src/series/implementations/dates_time.rs
@@ -11,6 +11,7 @@ use super::private;
 use super::IntoSeries;
 use super::SeriesTrait;
 use super::SeriesWrap;
+use super::*;
 use crate::chunked_array::{
     comparison::*,
     ops::{explode::ExplodeByOffsets, ToBitRepr},
@@ -226,6 +227,16 @@ macro_rules! impl_dyn_series {
         }
 
         impl SeriesTrait for SeriesWrap<$ca> {
+            fn is_sorted(&self) -> IsSorted {
+                if self.0.is_sorted() {
+                    IsSorted::Ascending
+                } else if self.0.is_sorted_reverse() {
+                    IsSorted::Descending
+                } else {
+                    IsSorted::Not
+                }
+            }
+
             #[cfg(feature = "interpolate")]
             fn interpolate(&self) -> Series {
                 self.0.interpolate().$into_logical().into_series()

diff --git a/polars/polars-core/src/series/implementations/datetime.rs b/polars/polars-core/src/series/implementations/datetime.rs
@@ -2,6 +2,7 @@ use super::private;
 use super::IntoSeries;
 use super::SeriesTrait;
 use super::SeriesWrap;
+use super::*;
 use crate::chunked_array::{
     comparison::*, ops::explode::ExplodeByOffsets, AsSinglePtr, ChunkIdIter,
 };
@@ -244,6 +245,16 @@ impl private::PrivateSeries for SeriesWrap<DatetimeChunked> {
 }
 
 impl SeriesTrait for SeriesWrap<DatetimeChunked> {
+    fn is_sorted(&self) -> IsSorted {
+        if self.0.is_sorted() {
+            IsSorted::Ascending
+        } else if self.0.is_sorted_reverse() {
+            IsSorted::Descending
+        } else {
+            IsSorted::Not
+        }
+    }
+
     #[cfg(feature = "interpolate")]
     fn interpolate(&self) -> Series {
         self.0

diff --git a/polars/polars-core/src/series/implementations/duration.rs b/polars/polars-core/src/series/implementations/duration.rs
@@ -2,6 +2,7 @@ use super::private;
 use super::IntoSeries;
 use super::SeriesTrait;
 use super::SeriesWrap;
+use super::*;
 use crate::chunked_array::{
     comparison::*, ops::explode::ExplodeByOffsets, AsSinglePtr, ChunkIdIter,
 };
@@ -236,6 +237,16 @@ impl private::PrivateSeries for SeriesWrap<DurationChunked> {
 }
 
 impl SeriesTrait for SeriesWrap<DurationChunked> {
+    fn is_sorted(&self) -> IsSorted {
+        if self.0.is_sorted() {
+            IsSorted::Ascending
+        } else if self.0.is_sorted_reverse() {
+            IsSorted::Descending
+        } else {
+            IsSorted::Not
+        }
+    }
+
     #[cfg(feature = "interpolate")]
     fn interpolate(&self) -> Series {
         self.0

diff --git a/polars/polars-core/src/series/implementations/floats.rs b/polars/polars-core/src/series/implementations/floats.rs
@@ -2,6 +2,7 @@ use super::private;
 use super::IntoSeries;
 use super::SeriesTrait;
 use super::SeriesWrap;
+use super::*;
 use crate::chunked_array::comparison::*;
 #[cfg(feature = "rolling_window")]
 use crate::chunked_array::ops::rolling_window::RollingOptions;
@@ -214,6 +215,16 @@ macro_rules! impl_dyn_series {
         }
 
         impl SeriesTrait for SeriesWrap<$ca> {
+            fn is_sorted(&self) -> IsSorted {
+                if self.0.is_sorted() {
+                    IsSorted::Ascending
+                } else if self.0.is_sorted_reverse() {
+                    IsSorted::Descending
+                } else {
+                    IsSorted::Not
+                }
+            }
+
             #[cfg(feature = "rolling_window")]
             fn rolling_apply(
                 &self,

diff --git a/polars/polars-core/src/series/implementations/mod.rs b/polars/polars-core/src/series/implementations/mod.rs
@@ -23,6 +23,7 @@ use std::any::Any;
 use super::private;
 use super::IntoSeries;
 use super::SeriesTrait;
+use super::*;
 use crate::chunked_array::comparison::*;
 #[cfg(feature = "rolling_window")]
 use crate::chunked_array::ops::rolling_window::RollingOptions;
@@ -260,6 +261,16 @@ macro_rules! impl_dyn_series {
         }
 
         impl SeriesTrait for SeriesWrap<$ca> {
+            fn is_sorted(&self) -> IsSorted {
+                if self.0.is_sorted() {
+                    IsSorted::Ascending
+                } else if self.0.is_sorted_reverse() {
+                    IsSorted::Descending
+                } else {
+                    IsSorted::Not
+                }
+            }
+
             #[cfg(feature = "rolling_window")]
             fn rolling_apply(
                 &self,

diff --git a/polars/polars-core/src/series/implementations/utf8.rs b/polars/polars-core/src/series/implementations/utf8.rs
@@ -1,6 +1,7 @@
 use super::private;
 use super::IntoSeries;
 use super::SeriesTrait;
+use super::*;
 use crate::chunked_array::comparison::*;
 use crate::chunked_array::{
     ops::{
@@ -109,6 +110,16 @@ impl private::PrivateSeries for SeriesWrap<Utf8Chunked> {
 }
 
 impl SeriesTrait for SeriesWrap<Utf8Chunked> {
+    fn is_sorted(&self) -> IsSorted {
+        if self.0.is_sorted() {
+            IsSorted::Ascending
+        } else if self.0.is_sorted_reverse() {
+            IsSorted::Descending
+        } else {
+            IsSorted::Not
+        }
+    }
+
     #[cfg(feature = "interpolate")]
     fn interpolate(&self) -> Series {
         self.0.clone().into_series()

diff --git a/polars/polars-core/src/series/mod.rs b/polars/polars-core/src/series/mod.rs
@@ -31,6 +31,8 @@ use std::hash::{Hash, Hasher};
 use std::ops::Deref;
 use std::sync::Arc;
 
+pub use series_trait::IsSorted;
+
 /// # Series
 /// The columnar data type for a DataFrame.
 ///

diff --git a/polars/polars-core/src/series/series_trait.rs b/polars/polars-core/src/series/series_trait.rs
@@ -12,6 +12,13 @@ use std::borrow::Cow;
 use std::ops::Deref;
 use std::sync::Arc;
 
+#[derive(Debug, Copy, Clone)]
+pub enum IsSorted {
+    Ascending,
+    Descending,
+    Not,
+}
+
 macro_rules! invalid_operation {
     ($s:expr) => {
         Err(PolarsError::InvalidOperation(
@@ -254,6 +261,11 @@ pub(crate) mod private {
 pub trait SeriesTrait:
     Send + Sync + private::PrivateSeries + private::PrivateSeriesNumeric
 {
+    /// Check if [`Series`] is sorted.
+    fn is_sorted(&self) -> IsSorted {
+        IsSorted::Not
+    }
+
     #[cfg(feature = "interpolate")]
     #[cfg_attr(docsrs, doc(cfg(feature = "interpolate")))]
     fn interpolate(&self) -> Series;

diff --git a/polars/polars-lazy/src/physical_plan/expressions/binary.rs b/polars/polars-lazy/src/physical_plan/expressions/binary.rs
@@ -103,7 +103,6 @@ impl PhysicalExpr for BinaryExpr {
             (AggState::AggregatedFlat(s), AggState::NotAggregated(_) | AggState::Literal(_))
                 if s.len() != df.height() =>
             {
-                dbg!("HIER");
                 // this is a flat series of len eq to group tuples
                 let l = ac_l.aggregated();
                 let l = l.as_ref();