test sortedness propagation (#3560)

pola-rs · Jun 2, 2022 · f1df2cb · f1df2cb
1 parent a98438b
commit f1df2cb
Show file tree

Hide file tree

Showing 11 changed files with 125 additions and 21 deletions.
diff --git a/polars/polars-core/src/frame/explode.rs b/polars/polars-core/src/frame/explode.rs
@@ -43,13 +43,21 @@ impl DataFrame {
                 .expect("cmp usize -> Ordering")
         });
 
+        let verbose = std::env::var("POLARS_VERBOSE").is_ok();
+
         // TODO: optimize this.
         // This is the slower easier option.
         // instead of filtering the whole dataframe first
         // drop empty list rows
         for col in &mut columns {
             if let Ok(ca) = col.list() {
                 if !ca.can_fast_explode() {
+                    if verbose {
+                        eprintln!(
+                            "could not fast explode column {}, running slower path",
+                            ca.name()
+                        );
+                    }
                     let (_, offsets) = get_exploded(col)?;
                     if offsets.is_empty() {
                         let mut out = self.slice(0, 0);
@@ -89,6 +97,8 @@ impl DataFrame {
                         .downcast_mut::<ListChunked>()
                         .unwrap();
                     ca.set_fast_explode();
+                } else if verbose {
+                    eprintln!("could fast explode column {}", ca.name());
                 }
             }
         }
@@ -107,7 +117,8 @@ impl DataFrame {
                 // expand all the other columns based the exploded first column
                 if i == 0 {
                     let row_idx = offsets_to_indexes(&offsets, exploded.len());
-                    let row_idx = IdxCa::from_vec("", row_idx);
+                    let mut row_idx = IdxCa::from_vec("", row_idx);
+                    row_idx.set_sorted(true);
                     // Safety
                     // We just created indices that are in bounds.
                     df = unsafe { df.take_unchecked(&row_idx) };

diff --git a/polars/polars-core/src/series/implementations/dates_time.rs b/polars/polars-core/src/series/implementations/dates_time.rs
@@ -326,9 +326,13 @@ macro_rules! impl_dyn_series {
             }
 
             unsafe fn take_unchecked(&self, idx: &IdxCa) -> Result<Series> {
-                Ok(ChunkTake::take_unchecked(self.0.deref(), idx.into())
-                    .$into_logical()
-                    .into_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())
+                }
+
+                Ok(out.$into_logical().into_series())
             }
 
             unsafe fn take_opt_iter_unchecked(&self, iter: &mut dyn TakeIteratorNulls) -> Series {

diff --git a/polars/polars-core/src/series/implementations/datetime.rs b/polars/polars-core/src/series/implementations/datetime.rs
@@ -333,7 +333,13 @@ impl SeriesTrait for SeriesWrap<DatetimeChunked> {
     }
 
     unsafe fn take_unchecked(&self, idx: &IdxCa) -> Result<Series> {
-        Ok(ChunkTake::take_unchecked(self.0.deref(), idx.into())
+        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())
+        }
+
+        Ok(out
             .into_datetime(self.0.time_unit(), self.0.time_zone().clone())
             .into_series())
     }

diff --git a/polars/polars-core/src/series/implementations/duration.rs b/polars/polars-core/src/series/implementations/duration.rs
@@ -360,9 +360,13 @@ impl SeriesTrait for SeriesWrap<DurationChunked> {
     }
 
     unsafe fn take_unchecked(&self, idx: &IdxCa) -> Result<Series> {
-        Ok(ChunkTake::take_unchecked(self.0.deref(), idx.into())
-            .into_duration(self.0.time_unit())
-            .into_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())
+        }
+
+        Ok(out.into_duration(self.0.time_unit()).into_series())
     }
 
     unsafe fn take_opt_iter_unchecked(&self, iter: &mut dyn TakeIteratorNulls) -> Series {

diff --git a/polars/polars-core/src/series/implementations/floats.rs b/polars/polars-core/src/series/implementations/floats.rs
@@ -337,7 +337,14 @@ macro_rules! impl_dyn_series {
                 } else {
                     Cow::Borrowed(idx)
                 };
-                Ok(ChunkTake::take_unchecked(&self.0, (&*idx).into()).into_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())
+                }
+
+                Ok(out.into_series())
             }
 
             unsafe fn take_opt_iter_unchecked(&self, iter: &mut dyn TakeIteratorNulls) -> Series {

diff --git a/polars/polars-core/src/series/implementations/mod.rs b/polars/polars-core/src/series/implementations/mod.rs
@@ -504,7 +504,11 @@ macro_rules! impl_dyn_series {
                 } else {
                     Cow::Borrowed(idx)
                 };
-                Ok(ChunkTake::take_unchecked(&self.0, (&*idx).into()).into_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())
+                }
+                Ok(out.into_series())
             }
 
             unsafe fn take_opt_iter_unchecked(&self, iter: &mut dyn TakeIteratorNulls) -> Series {

diff --git a/polars/polars-core/src/series/implementations/utf8.rs b/polars/polars-core/src/series/implementations/utf8.rs
@@ -221,7 +221,14 @@ impl SeriesTrait for SeriesWrap<Utf8Chunked> {
         } else {
             Cow::Borrowed(idx)
         };
-        Ok(ChunkTake::take_unchecked(&self.0, (&*idx).into()).into_series())
+
+        let mut out = ChunkTake::take_unchecked(&self.0, (&*idx).into()).into_series();
+
+        if self.0.is_sorted() && (idx.is_sorted() || idx.is_sorted_reverse()) {
+            out.set_sorted(idx.is_sorted_reverse())
+        }
+
+        Ok(out)
     }
 
     unsafe fn take_opt_iter_unchecked(&self, iter: &mut dyn TakeIteratorNulls) -> Series {

diff --git a/polars/polars-core/src/series/series_trait.rs b/polars/polars-core/src/series/series_trait.rs
@@ -9,7 +9,7 @@ use std::borrow::Cow;
 #[cfg(feature = "temporal")]
 use std::sync::Arc;
 
-#[derive(Debug, Copy, Clone)]
+#[derive(Debug, Copy, Clone, Eq, PartialEq)]
 pub enum IsSorted {
     Ascending,
     Descending,

diff --git a/polars/polars-lazy/src/frame/mod.rs b/polars/polars-lazy/src/frame/mod.rs
@@ -752,12 +752,17 @@ impl LazyFrame {
     ///        ])
     /// }
     /// ```
-    pub fn groupby<E: AsRef<[Expr]>>(self, by: E) -> LazyGroupBy {
+    pub fn groupby<E: AsRef<[IE]>, IE: Into<Expr> + Clone>(self, by: E) -> LazyGroupBy {
+        let keys = by
+            .as_ref()
+            .iter()
+            .map(|e| e.clone().into())
+            .collect::<Vec<_>>();
         let opt_state = self.get_opt_state();
         LazyGroupBy {
             logical_plan: self.logical_plan,
             opt_state,
-            keys: by.as_ref().to_vec(),
+            keys,
             maintain_order: false,
             dynamic_options: None,
             rolling_options: None,
@@ -797,12 +802,17 @@ impl LazyFrame {
     }
 
     /// Similar to groupby, but order of the DataFrame is maintained.
-    pub fn groupby_stable<E: AsRef<[Expr]>>(self, by: E) -> LazyGroupBy {
+    pub fn groupby_stable<E: AsRef<[IE]>, IE: Into<Expr> + Clone>(self, by: E) -> LazyGroupBy {
+        let keys = by
+            .as_ref()
+            .iter()
+            .map(|e| e.clone().into())
+            .collect::<Vec<_>>();
         let opt_state = self.get_opt_state();
         LazyGroupBy {
             logical_plan: self.logical_plan,
             opt_state,
-            keys: by.as_ref().to_vec(),
+            keys,
             maintain_order: true,
             dynamic_options: None,
             rolling_options: None,
@@ -983,8 +993,12 @@ impl LazyFrame {
     }
 
     /// Apply explode operation. [See eager explode](polars_core::frame::DataFrame::explode).
-    pub fn explode<E: AsRef<[Expr]>>(self, columns: E) -> LazyFrame {
-        let columns = columns.as_ref().to_vec();
+    pub fn explode<E: AsRef<[IE]>, IE: Into<Expr> + Clone>(self, columns: E) -> LazyFrame {
+        let columns = columns
+            .as_ref()
+            .iter()
+            .map(|e| e.clone().into())
+            .collect::<Vec<_>>();
         let opt_state = self.get_opt_state();
         let lp = self.get_plan_builder().explode(columns).build();
         Self::from_logical_plan(lp, opt_state)

diff --git a/polars/tests/it/lazy/queries.rs b/polars/tests/it/lazy/queries.rs
@@ -1,4 +1,5 @@
 use super::*;
+use polars_core::series::IsSorted;
 
 #[test]
 fn test_with_duplicate_column_empty_df() {
@@ -124,3 +125,26 @@ fn test_alias_before_cast() -> Result<()> {
     );
     Ok(())
 }
+
+#[test]
+fn test_sorted_path() -> Result<()> {
+    // start with a sorted column and see if the metadata remains preserved
+
+    let payloads = &[1, 2, 3];
+    let df = df![
+        "a"=> [AnyValue::List(Series::new("", payloads)), AnyValue::List(Series::new("", payloads)), AnyValue::List(Series::new("", payloads))]
+    ]?;
+
+    let out = df
+        .lazy()
+        .with_row_count("row_nr", None)
+        .explode(["a"])
+        .groupby(["row_nr"])
+        .agg([col("a").count().alias("count")])
+        .collect()?;
+
+    let s = out.column("row_nr")?;
+    assert_eq!(s.is_sorted(), IsSorted::Ascending);
+
+    Ok(())
+}
diff --git a/py-polars/src/series.rs b/py-polars/src/series.rs
@@ -260,9 +260,32 @@ impl PySeries {
     #[staticmethod]
     pub fn from_arrow(name: &str, array: &PyAny) -> PyResult<Self> {
         let arr = array_to_rust(array)?;
-        let series: Series =
-            std::convert::TryFrom::try_from((name, arr)).map_err(PyPolarsErr::from)?;
-        Ok(series.into())
+
+        match arr.data_type() {
+            ArrowDataType::LargeList(_) => {
+                let array = arr.as_any().downcast_ref::<LargeListArray>().unwrap();
+
+                let mut previous = 0;
+                let mut fast_explode = true;
+                for &o in array.offsets().as_slice()[1..].iter() {
+                    if o == previous {
+                        fast_explode = false;
+                        break;
+                    }
+                    previous = o;
+                }
+                let mut out = ListChunked::from_chunks(name, vec![arr]);
+                if fast_explode {
+                    out.set_fast_explode()
+                }
+                return Ok(out.into_series().into());
+            }
+            _ => {
+                let series: Series =
+                    std::convert::TryFrom::try_from((name, arr)).map_err(PyPolarsErr::from)?;
+                Ok(series.into())
+            }
+        }
     }
 
     #[staticmethod]