improve any/all perf

pola-rs · Jan 9, 2022 · cc659f8 · cc659f8
1 parent 1a74cbe
commit cc659f8
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 23 deletions.
diff --git a/polars/polars-core/src/chunked_array/comparison.rs b/polars/polars-core/src/chunked_array/comparison.rs
@@ -793,22 +793,14 @@ impl Not for BooleanChunked {
 }
 
 impl BooleanChunked {
-    /// Check if all values are true
-    pub fn all_true(&self) -> bool {
-        match self.sum() {
-            None => false,
-            Some(n) => (n as usize) == self.len(),
-        }
+    /// Check if all values are `true`
+    pub fn all(&self) -> bool {
+        self.downcast_iter().all(compute::boolean::all)
     }
-}
 
-impl BooleanChunked {
-    /// Check if all values are false
-    pub fn all_false(&self) -> bool {
-        match self.sum() {
-            None => false,
-            Some(n) => (n as usize) == 0,
-        }
+    /// Check if any value is `true`
+    pub fn any(&self) -> bool {
+        self.downcast_iter().any(compute::boolean::any)
     }
 }
 

diff --git a/polars/polars-core/src/frame/mod.rs b/polars/polars-core/src/frame/mod.rs
@@ -2611,7 +2611,7 @@ impl DataFrame {
     ///                         "ISIN" => &["US0378331005", "US5949181045"])?;
     /// let ca: ChunkedArray<BooleanType> = df.is_unique()?;
     ///
-    /// assert!(ca.all_true());
+    /// assert!(ca.all());
     /// # Ok::<(), PolarsError>(())
     /// ```
     pub fn is_unique(&self) -> Result<BooleanChunked> {
@@ -2630,7 +2630,7 @@ impl DataFrame {
     ///                         "ISIN" => &["US02079K3059", "US02079K1079"])?;
     /// let ca: ChunkedArray<BooleanType> = df.is_duplicated()?;
     ///
-    /// assert!(ca.all_false());
+    /// assert!(!ca.all());
     /// # Ok::<(), PolarsError>(())
     /// ```
     pub fn is_duplicated(&self) -> Result<BooleanChunked> {

diff --git a/polars/polars-lazy/src/dsl.rs b/polars/polars-lazy/src/dsl.rs
@@ -1990,11 +1990,10 @@ impl Expr {
         self.apply(
             move |s| {
                 let boolean = s.bool()?;
-                // TODO! Optimize this in arrow2/ polars-arrow
-                if boolean.all_false() {
-                    Ok(Series::new(s.name(), [false]))
-                } else {
+                if boolean.any() {
                     Ok(Series::new(s.name(), [true]))
+                } else {
+                    Ok(Series::new(s.name(), [false]))
                 }
             },
             GetOutput::from_type(DataType::Boolean),
@@ -2006,8 +2005,7 @@ impl Expr {
         self.apply(
             move |s| {
                 let boolean = s.bool()?;
-                // TODO! Optimize this in arrow2/ polars-arrow
-                if boolean.all_true() {
+                if boolean.all() {
                     Ok(Series::new(s.name(), [true]))
                 } else {
                     Ok(Series::new(s.name(), [false]))

diff --git a/polars/polars-lazy/src/tests/queries.rs b/polars/polars-lazy/src/tests/queries.rs
@@ -55,7 +55,7 @@ fn test_lazy_exec() {
         .unwrap();
 
     let check = new.column("sepal.width").unwrap().f64().unwrap().gt(3.4);
-    assert!(check.all_true())
+    assert!(check.all())
 }
 
 #[test]