perf: Use downcast_ref instead of dtype equality in `<dyn SeriesTrait…

… as AsRef<ChunkedArray<T>>` (#20664)
pola-rs · Jan 10, 2025 · 94b4087 · 94b4087
1 parent b816b97
commit 94b4087
Show file tree

Hide file tree

Showing 21 changed files with 118 additions and 61 deletions.
diff --git a/crates/polars-core/src/datatypes/mod.rs b/crates/polars-core/src/datatypes/mod.rs
@@ -78,6 +78,7 @@ pub unsafe trait PolarsDataType: Send + Sync + Sized + 'static {
     type HasViews;
     type IsStruct;
     type IsObject;
+    type IsLogical;
 
     fn get_dtype() -> DataType
     where
@@ -95,6 +96,7 @@ where
         HasViews = FalseT,
         IsStruct = FalseT,
         IsObject = FalseT,
+        IsLogical = FalseT,
     >,
 {
     type Native: NumericNative;
@@ -117,6 +119,7 @@ macro_rules! impl_polars_num_datatype {
             type HasViews = FalseT;
             type IsStruct = FalseT;
             type IsObject = FalseT;
+            type IsLogical = FalseT;
 
             #[inline]
             fn get_dtype() -> DataType {
@@ -133,7 +136,7 @@ macro_rules! impl_polars_num_datatype {
 }
 
 macro_rules! impl_polars_datatype_pass_dtype {
-    ($ca:ident, $dtype:expr, $arr:ty, $lt:lifetime, $phys:ty, $zerophys:ty, $owned_phys:ty, $has_views:ident) => {
+    ($ca:ident, $dtype:expr, $arr:ty, $lt:lifetime, $phys:ty, $zerophys:ty, $owned_phys:ty, $has_views:ident, $is_logical:ident) => {
         #[derive(Clone, Copy)]
         pub struct $ca {}
 
@@ -146,6 +149,7 @@ macro_rules! impl_polars_datatype_pass_dtype {
             type HasViews = $has_views;
             type IsStruct = FalseT;
             type IsObject = FalseT;
+            type IsLogical = $is_logical;
 
             #[inline]
             fn get_dtype() -> DataType {
@@ -164,13 +168,14 @@ macro_rules! impl_polars_binview_datatype {
             $phys,
             $zerophys,
             $owned_phys,
-            TrueT
+            TrueT,
+            FalseT
         );
     };
 }
 
 macro_rules! impl_polars_datatype {
-    ($ca:ident, $variant:ident, $arr:ty, $lt:lifetime, $phys:ty, $zerophys:ty, $owned_phys:ty) => {
+    ($ca:ident, $variant:ident, $arr:ty, $lt:lifetime, $phys:ty, $zerophys:ty, $owned_phys:ty, $is_logical:ident) => {
         impl_polars_datatype_pass_dtype!(
             $ca,
             DataType::$variant,
@@ -179,7 +184,8 @@ macro_rules! impl_polars_datatype {
             $phys,
             $zerophys,
             $owned_phys,
-            FalseT
+            FalseT,
+            $is_logical
         );
     };
 }
@@ -197,18 +203,18 @@ impl_polars_num_datatype!(PolarsIntegerType, Int64Type, Int64, i64, i64);
 impl_polars_num_datatype!(PolarsIntegerType, Int128Type, Int128, i128, i128);
 impl_polars_num_datatype!(PolarsFloatType, Float32Type, Float32, f32, f32);
 impl_polars_num_datatype!(PolarsFloatType, Float64Type, Float64, f64, f64);
-impl_polars_datatype!(DateType, Date, PrimitiveArray<i32>, 'a, i32, i32, i32);
-impl_polars_datatype!(TimeType, Time, PrimitiveArray<i64>, 'a, i64, i64, i64);
+impl_polars_datatype!(DateType, Date, PrimitiveArray<i32>, 'a, i32, i32, i32, TrueT);
+impl_polars_datatype!(TimeType, Time, PrimitiveArray<i64>, 'a, i64, i64, i64, TrueT);
 impl_polars_binview_datatype!(StringType, String, Utf8ViewArray, 'a, &'a str, Option<&'a str>, String);
 impl_polars_binview_datatype!(BinaryType, Binary, BinaryViewArray, 'a, &'a [u8], Option<&'a [u8]>, Box<[u8]>);
-impl_polars_datatype!(BinaryOffsetType, BinaryOffset, BinaryArray<i64>, 'a, &'a [u8], Option<&'a [u8]>, Box<[u8]>);
-impl_polars_datatype!(BooleanType, Boolean, BooleanArray, 'a, bool, bool, bool);
+impl_polars_datatype!(BinaryOffsetType, BinaryOffset, BinaryArray<i64>, 'a, &'a [u8], Option<&'a [u8]>, Box<[u8]>, FalseT);
+impl_polars_datatype!(BooleanType, Boolean, BooleanArray, 'a, bool, bool, bool, FalseT);
 
 #[cfg(feature = "dtype-decimal")]
-impl_polars_datatype_pass_dtype!(DecimalType, DataType::Unknown(UnknownKind::Any), PrimitiveArray<i128>, 'a, i128, i128, i128, FalseT);
-impl_polars_datatype_pass_dtype!(DatetimeType, DataType::Unknown(UnknownKind::Any), PrimitiveArray<i64>, 'a, i64, i64, i64, FalseT);
-impl_polars_datatype_pass_dtype!(DurationType, DataType::Unknown(UnknownKind::Any), PrimitiveArray<i64>, 'a, i64, i64, i64, FalseT);
-impl_polars_datatype_pass_dtype!(CategoricalType, DataType::Unknown(UnknownKind::Any), PrimitiveArray<u32>, 'a, u32, u32, u32, FalseT);
+impl_polars_datatype_pass_dtype!(DecimalType, DataType::Unknown(UnknownKind::Any), PrimitiveArray<i128>, 'a, i128, i128, i128, FalseT, TrueT);
+impl_polars_datatype_pass_dtype!(DatetimeType, DataType::Unknown(UnknownKind::Any), PrimitiveArray<i64>, 'a, i64, i64, i64, FalseT, TrueT);
+impl_polars_datatype_pass_dtype!(DurationType, DataType::Unknown(UnknownKind::Any), PrimitiveArray<i64>, 'a, i64, i64, i64, FalseT, TrueT);
+impl_polars_datatype_pass_dtype!(CategoricalType, DataType::Unknown(UnknownKind::Any), PrimitiveArray<u32>, 'a, u32, u32, u32, FalseT, TrueT);
 
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 pub struct ListType {}
@@ -221,6 +227,7 @@ unsafe impl PolarsDataType for ListType {
     type HasViews = FalseT;
     type IsStruct = FalseT;
     type IsObject = FalseT;
+    type IsLogical = FalseT;
 
     fn get_dtype() -> DataType {
         // Null as we cannot know anything without self.
@@ -245,6 +252,7 @@ unsafe impl PolarsDataType for StructType {
     type HasViews = FalseT;
     type IsStruct = TrueT;
     type IsObject = FalseT;
+    type IsLogical = FalseT;
 
     fn get_dtype() -> DataType
     where
@@ -266,6 +274,7 @@ unsafe impl PolarsDataType for FixedSizeListType {
     type HasViews = FalseT;
     type IsStruct = FalseT;
     type IsObject = FalseT;
+    type IsLogical = FalseT;
 
     fn get_dtype() -> DataType {
         // Null as we cannot know anything without self.
@@ -285,6 +294,7 @@ unsafe impl<T: PolarsObject> PolarsDataType for ObjectType<T> {
     type HasViews = FalseT;
     type IsStruct = FalseT;
     type IsObject = TrueT;
+    type IsLogical = FalseT;
 
     fn get_dtype() -> DataType {
         DataType::Object(T::type_name(), None)

diff --git a/crates/polars-core/src/series/implementations/array.rs b/crates/polars-core/src/series/implementations/array.rs
@@ -207,12 +207,11 @@ impl SeriesTrait for SeriesWrap<ArrayChunked> {
     fn clone_inner(&self) -> Arc<dyn SeriesTrait> {
         Arc::new(SeriesWrap(Clone::clone(&self.0)))
     }
+
     fn as_any(&self) -> &dyn Any {
         &self.0
     }
 
-    /// Get a hold to self as `Any` trait reference.
-    /// Only implemented for ObjectType
     fn as_any_mut(&mut self) -> &mut dyn Any {
         &mut self.0
     }

diff --git a/crates/polars-core/src/series/implementations/binary.rs b/crates/polars-core/src/series/implementations/binary.rs
@@ -247,4 +247,8 @@ impl SeriesTrait for SeriesWrap<BinaryChunked> {
     fn as_any(&self) -> &dyn Any {
         &self.0
     }
+
+    fn as_any_mut(&mut self) -> &mut dyn Any {
+        &mut self.0
+    }
 }
diff --git a/crates/polars-core/src/series/implementations/binary_offset.rs b/crates/polars-core/src/series/implementations/binary_offset.rs
@@ -189,4 +189,8 @@ impl SeriesTrait for SeriesWrap<BinaryOffsetChunked> {
     fn as_any(&self) -> &dyn Any {
         &self.0
     }
+
+    fn as_any_mut(&mut self) -> &mut dyn Any {
+        &mut self.0
+    }
 }
diff --git a/crates/polars-core/src/series/implementations/boolean.rs b/crates/polars-core/src/series/implementations/boolean.rs
@@ -354,4 +354,8 @@ impl SeriesTrait for SeriesWrap<BooleanChunked> {
     fn as_any(&self) -> &dyn Any {
         &self.0
     }
+
+    fn as_any_mut(&mut self) -> &mut dyn Any {
+        &mut self.0
+    }
 }
diff --git a/crates/polars-core/src/series/implementations/categorical.rs b/crates/polars-core/src/series/implementations/categorical.rs
@@ -293,9 +293,14 @@ impl SeriesTrait for SeriesWrap<CategoricalChunked> {
     fn max_reduce(&self) -> PolarsResult<Scalar> {
         Ok(ChunkAggSeries::max_reduce(&self.0))
     }
+
     fn as_any(&self) -> &dyn Any {
         &self.0
     }
+
+    fn as_any_mut(&mut self) -> &mut dyn Any {
+        &mut self.0
+    }
 }
 
 impl private::PrivateSeriesNumeric for SeriesWrap<CategoricalChunked> {

diff --git a/crates/polars-core/src/series/implementations/date.rs b/crates/polars-core/src/series/implementations/date.rs
@@ -351,6 +351,10 @@ impl SeriesTrait for SeriesWrap<DateChunked> {
     fn as_any(&self) -> &dyn Any {
         &self.0
     }
+
+    fn as_any_mut(&mut self) -> &mut dyn Any {
+        &mut self.0
+    }
 }
 
 impl private::PrivateSeriesNumeric for SeriesWrap<DateChunked> {

diff --git a/crates/polars-core/src/series/implementations/datetime.rs b/crates/polars-core/src/series/implementations/datetime.rs
@@ -360,7 +360,12 @@ impl SeriesTrait for SeriesWrap<DatetimeChunked> {
     fn clone_inner(&self) -> Arc<dyn SeriesTrait> {
         Arc::new(SeriesWrap(Clone::clone(&self.0)))
     }
+
     fn as_any(&self) -> &dyn Any {
         &self.0
     }
+
+    fn as_any_mut(&mut self) -> &mut dyn Any {
+        &mut self.0
+    }
 }
diff --git a/crates/polars-core/src/series/implementations/decimal.rs b/crates/polars-core/src/series/implementations/decimal.rs
@@ -419,4 +419,8 @@ impl SeriesTrait for SeriesWrap<DecimalChunked> {
     fn as_any(&self) -> &dyn Any {
         &self.0
     }
+
+    fn as_any_mut(&mut self) -> &mut dyn Any {
+        &mut self.0
+    }
 }
diff --git a/crates/polars-core/src/series/implementations/duration.rs b/crates/polars-core/src/series/implementations/duration.rs
@@ -520,4 +520,8 @@ impl SeriesTrait for SeriesWrap<DurationChunked> {
     fn as_any(&self) -> &dyn Any {
         &self.0
     }
+
+    fn as_any_mut(&mut self) -> &mut dyn Any {
+        &mut self.0
+    }
 }
diff --git a/crates/polars-core/src/series/implementations/floats.rs b/crates/polars-core/src/series/implementations/floats.rs
@@ -362,9 +362,14 @@ macro_rules! impl_dyn_series {
             fn checked_div(&self, rhs: &Series) -> PolarsResult<Series> {
                 self.0.checked_div(rhs)
             }
+
             fn as_any(&self) -> &dyn Any {
                 &self.0
             }
+
+            fn as_any_mut(&mut self) -> &mut dyn Any {
+                &mut self.0
+            }
         }
     };
 }

diff --git a/crates/polars-core/src/series/implementations/list.rs b/crates/polars-core/src/series/implementations/list.rs
@@ -254,12 +254,11 @@ impl SeriesTrait for SeriesWrap<ListChunked> {
     fn clone_inner(&self) -> Arc<dyn SeriesTrait> {
         Arc::new(SeriesWrap(Clone::clone(&self.0)))
     }
+
     fn as_any(&self) -> &dyn Any {
         &self.0
     }
 
-    /// Get a hold to self as `Any` trait reference.
-    /// Only implemented for ObjectType
     fn as_any_mut(&mut self) -> &mut dyn Any {
         &mut self.0
     }

diff --git a/crates/polars-core/src/series/implementations/mod.rs b/crates/polars-core/src/series/implementations/mod.rs
@@ -439,6 +439,10 @@ macro_rules! impl_dyn_series {
             fn as_any(&self) -> &dyn Any {
                 &self.0
             }
+
+            fn as_any_mut(&mut self) -> &mut dyn Any {
+                &mut self.0
+            }
         }
     };
 }

diff --git a/crates/polars-core/src/series/implementations/null.rs b/crates/polars-core/src/series/implementations/null.rs
@@ -316,9 +316,14 @@ impl SeriesTrait for NullChunked {
     fn clone_inner(&self) -> Arc<dyn SeriesTrait> {
         Arc::new(self.clone())
     }
+
     fn as_any(&self) -> &dyn Any {
         self
     }
+
+    fn as_any_mut(&mut self) -> &mut dyn Any {
+        self
+    }
 }
 
 unsafe impl IntoSeries for NullChunked {

diff --git a/crates/polars-core/src/series/implementations/object.rs b/crates/polars-core/src/series/implementations/object.rs
@@ -237,6 +237,10 @@ where
     fn as_any(&self) -> &dyn Any {
         &self.0
     }
+
+    fn as_any_mut(&mut self) -> &mut dyn Any {
+        &mut self.0
+    }
 }
 
 #[cfg(test)]

diff --git a/crates/polars-core/src/series/implementations/string.rs b/crates/polars-core/src/series/implementations/string.rs
@@ -261,4 +261,8 @@ impl SeriesTrait for SeriesWrap<StringChunked> {
     fn as_any(&self) -> &dyn Any {
         &self.0
     }
+
+    fn as_any_mut(&mut self) -> &mut dyn Any {
+        &mut self.0
+    }
 }
diff --git a/crates/polars-core/src/series/implementations/struct_.rs b/crates/polars-core/src/series/implementations/struct_.rs
@@ -257,6 +257,10 @@ impl SeriesTrait for SeriesWrap<StructChunked> {
         &self.0
     }
 
+    fn as_any_mut(&mut self) -> &mut dyn Any {
+        &mut self.0
+    }
+
     fn sort_with(&self, options: SortOptions) -> PolarsResult<Series> {
         Ok(self.0.sort_with(options).into_series())
     }

diff --git a/crates/polars-core/src/series/implementations/time.rs b/crates/polars-core/src/series/implementations/time.rs
@@ -324,6 +324,10 @@ impl SeriesTrait for SeriesWrap<TimeChunked> {
     fn as_any(&self) -> &dyn Any {
         &self.0
     }
+
+    fn as_any_mut(&mut self) -> &mut dyn Any {
+        &mut self.0
+    }
 }
 
 impl private::PrivateSeriesNumeric for SeriesWrap<TimeChunked> {
-Original file line number
+Diff line change
@@ Expand Up / @@ -439,6 +439,10 @@ macro_rules! impl_dyn_series { @@
                 fn as_any(&self) -> &dyn Any {
                     &self.0
                 }
+                fn as_any_mut(&mut self) -> &mut dyn Any {
+                    &mut self.0
+                }
             }
         };
     }
@@ Expand Down @@