refactor: genericize PolarsDataType (pola-rs#10952)

sd2k · Sep 7, 2023 · 64bf7bf · 64bf7bf
1 parent faa71c5
commit 64bf7bf
Show file tree

Hide file tree

Showing 11 changed files with 157 additions and 302 deletions.
diff --git a/crates/polars-core/src/chunked_array/from.rs b/crates/polars-core/src/chunked_array/from.rs
@@ -72,8 +72,8 @@ fn from_chunks_list_dtype(chunks: &mut Vec<ArrayRef>, dtype: DataType) -> DataTy
 
 impl<T, A> From<A> for ChunkedArray<T>
 where
-    T: PolarsDataType,
-    A: StaticallyMatchesPolarsType<T> + Array,
+    T: PolarsDataType<Array = A>,
+    A: Array,
 {
     fn from(arr: A) -> Self {
         Self::with_chunk("", arr)
@@ -86,15 +86,17 @@ where
 {
     pub fn with_chunk<A>(name: &str, arr: A) -> Self
     where
-        A: StaticallyMatchesPolarsType<T> + Array,
+        A: Array,
+        T: PolarsDataType<Array = A>,
     {
         unsafe { Self::from_chunks(name, vec![Box::new(arr)]) }
     }
 
     pub fn from_chunk_iter<I>(name: &str, iter: I) -> Self
     where
         I: IntoIterator,
-        <I as IntoIterator>::Item: StaticallyMatchesPolarsType<T> + Array,
+        T: PolarsDataType<Array = <I as IntoIterator>::Item>,
+        <I as IntoIterator>::Item: Array,
     {
         let chunks = iter
             .into_iter()
@@ -106,7 +108,8 @@ where
     pub fn try_from_chunk_iter<I, A, E>(name: &str, iter: I) -> Result<Self, E>
     where
         I: IntoIterator<Item = Result<A, E>>,
-        A: StaticallyMatchesPolarsType<T> + Array,
+        T: PolarsDataType<Array = A>,
+        A: Array,
     {
         let chunks: Result<_, _> = iter
             .into_iter()

diff --git a/crates/polars-core/src/chunked_array/list/iterator.rs b/crates/polars-core/src/chunked_array/list/iterator.rs
@@ -173,8 +173,7 @@ impl ListChunked {
     where
         V: PolarsDataType,
         F: FnMut(Option<UnstableSeries<'a>>) -> Option<K> + Copy,
-        K: ArrayFromElementIter,
-        K::ArrayType: StaticallyMatchesPolarsType<V>,
+        K: ArrayFromElementIter<ArrayType = V::Array>,
     {
         // TODO! make an amortized iter that does not flatten
 

diff --git a/crates/polars-core/src/chunked_array/ops/apply.rs b/crates/polars-core/src/chunked_array/ops/apply.rs
@@ -17,14 +17,12 @@ use crate::utils::CustomIterTools;
 impl<T> ChunkedArray<T>
 where
     T: PolarsDataType,
-    Self: HasUnderlyingArray,
 {
     pub fn apply_values_generic<'a, U, K, F>(&'a self, op: F) -> ChunkedArray<U>
     where
         U: PolarsDataType,
-        F: FnMut(<<Self as HasUnderlyingArray>::ArrayT as StaticArray>::ValueT<'a>) -> K + Copy,
-        K: ArrayFromElementIter,
-        K::ArrayType: StaticallyMatchesPolarsType<U>,
+        F: FnMut(T::Physical<'a>) -> K + Copy,
+        K: ArrayFromElementIter<ArrayType = U::Array>,
     {
         let iter = self.downcast_iter().map(|arr| {
             let element_iter = arr.values_iter().map(op);
@@ -38,10 +36,8 @@ where
     pub fn try_apply_values_generic<'a, U, K, F, E>(&'a self, op: F) -> Result<ChunkedArray<U>, E>
     where
         U: PolarsDataType,
-        F: FnMut(<<Self as HasUnderlyingArray>::ArrayT as StaticArray>::ValueT<'a>) -> Result<K, E>
-            + Copy,
-        K: ArrayFromElementIter,
-        K::ArrayType: StaticallyMatchesPolarsType<U>,
+        F: FnMut(T::Physical<'a>) -> Result<K, E> + Copy,
+        K: ArrayFromElementIter<ArrayType = U::Array>,
         E: Error,
     {
         let iter = self.downcast_iter().map(|arr| {
@@ -56,12 +52,8 @@ where
     pub fn try_apply_generic<'a, U, K, F, E>(&'a self, op: F) -> Result<ChunkedArray<U>, E>
     where
         U: PolarsDataType,
-        F: FnMut(
-                Option<<<Self as HasUnderlyingArray>::ArrayT as StaticArray>::ValueT<'a>>,
-            ) -> Result<Option<K>, E>
-            + Copy,
-        K: ArrayFromElementIter,
-        K::ArrayType: StaticallyMatchesPolarsType<U>,
+        F: FnMut(Option<T::Physical<'a>>) -> Result<Option<K>, E> + Copy,
+        K: ArrayFromElementIter<ArrayType = U::Array>,
         E: Error,
     {
         let iter = self.downcast_iter().map(|arr| {
@@ -76,11 +68,8 @@ where
     pub fn apply_generic<'a, U, K, F>(&'a self, mut op: F) -> ChunkedArray<U>
     where
         U: PolarsDataType,
-        F: FnMut(
-            Option<<<Self as HasUnderlyingArray>::ArrayT as StaticArray>::ValueT<'a>>,
-        ) -> Option<K>,
-        K: ArrayFromElementIter,
-        K::ArrayType: StaticallyMatchesPolarsType<U>,
+        F: FnMut(Option<T::Physical<'a>>) -> Option<K>,
+        K: ArrayFromElementIter<ArrayType = U::Array>,
     {
         if self.null_count() == 0 {
             let iter = self.downcast_iter().map(|arr| {

diff --git a/crates/polars-core/src/chunked_array/ops/arity.rs b/crates/polars-core/src/chunked_array/ops/arity.rs
@@ -3,10 +3,7 @@ use std::error::Error;
 use arrow::array::Array;
 use polars_arrow::utils::combine_validities_and;
 
-use crate::datatypes::{
-    ArrayFromElementIter, HasUnderlyingArray, PolarsNumericType, StaticArray,
-    StaticallyMatchesPolarsType,
-};
+use crate::datatypes::{ArrayFromElementIter, PolarsNumericType, StaticArray};
 use crate::prelude::{ChunkedArray, PolarsDataType};
 use crate::utils::align_chunks_binary;
 
@@ -20,14 +17,8 @@ where
     T: PolarsDataType,
     U: PolarsDataType,
     V: PolarsDataType,
-    ChunkedArray<T>: HasUnderlyingArray,
-    ChunkedArray<U>: HasUnderlyingArray,
-    F: for<'a> FnMut(
-        Option<<<ChunkedArray<T> as HasUnderlyingArray>::ArrayT as StaticArray>::ValueT<'a>>,
-        Option<<<ChunkedArray<U> as HasUnderlyingArray>::ArrayT as StaticArray>::ValueT<'a>>,
-    ) -> Option<K>,
-    K: ArrayFromElementIter,
-    K::ArrayType: StaticallyMatchesPolarsType<V>,
+    F: for<'a> FnMut(Option<T::Physical<'a>>, Option<U::Physical<'a>>) -> Option<K>,
+    K: ArrayFromElementIter<ArrayType = V::Array>,
 {
     let (lhs, rhs) = align_chunks_binary(lhs, rhs);
     let iter = lhs
@@ -53,14 +44,8 @@ where
     T: PolarsDataType,
     U: PolarsDataType,
     V: PolarsDataType,
-    ChunkedArray<T>: HasUnderlyingArray,
-    ChunkedArray<U>: HasUnderlyingArray,
-    F: for<'a> FnMut(
-        Option<<<ChunkedArray<T> as HasUnderlyingArray>::ArrayT as StaticArray>::ValueT<'a>>,
-        Option<<<ChunkedArray<U> as HasUnderlyingArray>::ArrayT as StaticArray>::ValueT<'a>>,
-    ) -> Result<Option<K>, E>,
-    K: ArrayFromElementIter,
-    K::ArrayType: StaticallyMatchesPolarsType<V>,
+    F: for<'a> FnMut(Option<T::Physical<'a>>, Option<U::Physical<'a>>) -> Result<Option<K>, E>,
+    K: ArrayFromElementIter<ArrayType = V::Array>,
     E: Error,
 {
     let (lhs, rhs) = align_chunks_binary(lhs, rhs);
@@ -87,14 +72,8 @@ where
     T: PolarsDataType,
     U: PolarsDataType,
     V: PolarsNumericType,
-    ChunkedArray<T>: HasUnderlyingArray,
-    ChunkedArray<U>: HasUnderlyingArray,
-    F: for<'a> FnMut(
-        <<ChunkedArray<T> as HasUnderlyingArray>::ArrayT as StaticArray>::ValueT<'a>,
-        <<ChunkedArray<U> as HasUnderlyingArray>::ArrayT as StaticArray>::ValueT<'a>,
-    ) -> K,
-    K: ArrayFromElementIter,
-    K::ArrayType: StaticallyMatchesPolarsType<V>,
+    F: for<'a> FnMut(T::Physical<'a>, U::Physical<'a>) -> K,
+    K: ArrayFromElementIter<ArrayType = V::Array>,
 {
     let (lhs, rhs) = align_chunks_binary(lhs, rhs);
     let iter = lhs
@@ -124,14 +103,8 @@ where
     T: PolarsDataType,
     U: PolarsDataType,
     V: PolarsNumericType,
-    ChunkedArray<T>: HasUnderlyingArray,
-    ChunkedArray<U>: HasUnderlyingArray,
-    F: for<'a> FnMut(
-        <<ChunkedArray<T> as HasUnderlyingArray>::ArrayT as StaticArray>::ValueT<'a>,
-        <<ChunkedArray<U> as HasUnderlyingArray>::ArrayT as StaticArray>::ValueT<'a>,
-    ) -> Result<K, E>,
-    K: ArrayFromElementIter,
-    K::ArrayType: StaticallyMatchesPolarsType<V>,
+    F: for<'a> FnMut(T::Physical<'a>, U::Physical<'a>) -> Result<K, E>,
+    K: ArrayFromElementIter<ArrayType = V::Array>,
     E: Error,
 {
     let (lhs, rhs) = align_chunks_binary(lhs, rhs);
@@ -163,14 +136,9 @@ pub fn binary_mut_with_options<T, U, V, F, Arr>(
 where
     T: PolarsDataType,
     U: PolarsDataType,
-    V: PolarsDataType,
-    ChunkedArray<T>: HasUnderlyingArray,
-    ChunkedArray<U>: HasUnderlyingArray,
-    Arr: Array + StaticallyMatchesPolarsType<V>,
-    F: FnMut(
-        &<ChunkedArray<T> as HasUnderlyingArray>::ArrayT,
-        &<ChunkedArray<U> as HasUnderlyingArray>::ArrayT,
-    ) -> Arr,
+    V: PolarsDataType<Array = Arr>,
+    Arr: Array,
+    F: FnMut(&T::Array, &U::Array) -> Arr,
 {
     let (lhs, rhs) = align_chunks_binary(lhs, rhs);
     let iter = lhs
@@ -189,14 +157,9 @@ pub fn binary<T, U, V, F, Arr>(
 where
     T: PolarsDataType,
     U: PolarsDataType,
-    V: PolarsDataType,
-    ChunkedArray<T>: HasUnderlyingArray,
-    ChunkedArray<U>: HasUnderlyingArray,
-    Arr: Array + StaticallyMatchesPolarsType<V>,
-    F: FnMut(
-        &<ChunkedArray<T> as HasUnderlyingArray>::ArrayT,
-        &<ChunkedArray<U> as HasUnderlyingArray>::ArrayT,
-    ) -> Arr,
+    V: PolarsDataType<Array = Arr>,
+    Arr: Array,
+    F: FnMut(&T::Array, &U::Array) -> Arr,
 {
     binary_mut_with_options(lhs, rhs, op, lhs.name())
 }
@@ -210,14 +173,9 @@ pub fn try_binary<T, U, V, F, Arr, E>(
 where
     T: PolarsDataType,
     U: PolarsDataType,
-    V: PolarsDataType,
-    ChunkedArray<T>: HasUnderlyingArray,
-    ChunkedArray<U>: HasUnderlyingArray,
-    Arr: Array + StaticallyMatchesPolarsType<V>,
-    F: FnMut(
-        &<ChunkedArray<T> as HasUnderlyingArray>::ArrayT,
-        &<ChunkedArray<U> as HasUnderlyingArray>::ArrayT,
-    ) -> Result<Arr, E>,
+    V: PolarsDataType<Array = Arr>,
+    Arr: Array,
+    F: FnMut(&T::Array, &U::Array) -> Result<Arr, E>,
     E: Error,
 {
     let (lhs, rhs) = align_chunks_binary(lhs, rhs);
@@ -243,12 +201,7 @@ pub unsafe fn binary_unchecked_same_type<T, U, F>(
 where
     T: PolarsDataType,
     U: PolarsDataType,
-    ChunkedArray<T>: HasUnderlyingArray,
-    ChunkedArray<U>: HasUnderlyingArray,
-    F: FnMut(
-        &<ChunkedArray<T> as HasUnderlyingArray>::ArrayT,
-        &<ChunkedArray<U> as HasUnderlyingArray>::ArrayT,
-    ) -> Box<dyn Array>,
+    F: FnMut(&T::Array, &U::Array) -> Box<dyn Array>,
 {
     let (lhs, rhs) = align_chunks_binary(lhs, rhs);
     let chunks = lhs
@@ -274,12 +227,7 @@ pub unsafe fn try_binary_unchecked_same_type<T, U, F, E>(
 where
     T: PolarsDataType,
     U: PolarsDataType,
-    ChunkedArray<T>: HasUnderlyingArray,
-    ChunkedArray<U>: HasUnderlyingArray,
-    F: FnMut(
-        &<ChunkedArray<T> as HasUnderlyingArray>::ArrayT,
-        &<ChunkedArray<U> as HasUnderlyingArray>::ArrayT,
-    ) -> Result<Box<dyn Array>, E>,
+    F: FnMut(&T::Array, &U::Array) -> Result<Box<dyn Array>, E>,
     E: Error,
 {
     let (lhs, rhs) = align_chunks_binary(lhs, rhs);

diff --git a/crates/polars-core/src/chunked_array/ops/downcast.rs b/crates/polars-core/src/chunked_array/ops/downcast.rs
@@ -47,18 +47,13 @@ impl<'a, T> Chunks<'a, T> {
 }
 
 #[doc(hidden)]
-impl<T: PolarsDataType> ChunkedArray<T>
-where
-    Self: HasUnderlyingArray,
-{
+impl<T: PolarsDataType> ChunkedArray<T> {
     #[inline]
-    pub fn downcast_iter(
-        &self,
-    ) -> impl Iterator<Item = &<Self as HasUnderlyingArray>::ArrayT> + DoubleEndedIterator {
+    pub fn downcast_iter(&self) -> impl Iterator<Item = &T::Array> + DoubleEndedIterator {
         self.chunks.iter().map(|arr| {
-            // SAFETY: HasUnderlyingArray guarantees this is correct.
+            // SAFETY: T::Array guarantees this is correct.
             let arr = &**arr;
-            unsafe { &*(arr as *const dyn Array as *const <Self as HasUnderlyingArray>::ArrayT) }
+            unsafe { &*(arr as *const dyn Array as *const T::Array) }
         })
     }
 
@@ -69,16 +64,16 @@ where
     #[inline]
     pub unsafe fn downcast_iter_mut(
         &mut self,
-    ) -> impl Iterator<Item = &mut <Self as HasUnderlyingArray>::ArrayT> + DoubleEndedIterator {
+    ) -> impl Iterator<Item = &mut T::Array> + DoubleEndedIterator {
         self.chunks.iter_mut().map(|arr| {
-            // SAFETY: HasUnderlyingArray guarantees this is correct.
+            // SAFETY: T::Array guarantees this is correct.
             let arr = &mut **arr;
-            &mut *(arr as *mut dyn Array as *mut <Self as HasUnderlyingArray>::ArrayT)
+            &mut *(arr as *mut dyn Array as *mut T::Array)
         })
     }
 
     #[inline]
-    pub fn downcast_chunks(&self) -> Chunks<'_, <Self as HasUnderlyingArray>::ArrayT> {
+    pub fn downcast_chunks(&self) -> Chunks<'_, T::Array> {
         Chunks::new(&self.chunks)
     }
 

diff --git a/crates/polars-core/src/chunked_array/ops/mod.rs b/crates/polars-core/src/chunked_array/ops/mod.rs
@@ -589,10 +589,9 @@ macro_rules! impl_chunk_expand {
     }};
 }
 
-impl<T: PolarsDataType> ChunkExpandAtIndex<T> for ChunkedArray<T>
+impl<T: PolarsNumericType> ChunkExpandAtIndex<T> for ChunkedArray<T>
 where
     ChunkedArray<T>: ChunkFull<T::Native> + TakeRandom<Item = T::Native>,
-    T: PolarsNumericType,
 {
     fn new_from_index(&self, index: usize, length: usize) -> ChunkedArray<T> {
         let mut out = impl_chunk_expand!(self, length, index);