refactor: Add a Column::Partitioned variant (#19557)

pola-rs · Nov 1, 2024 · 78f56c2 · 78f56c2
1 parent 6abcc54
commit 78f56c2
Show file tree

Hide file tree

Showing 8 changed files with 383 additions and 48 deletions.
diff --git a/crates/polars-core/src/datatypes/any_value.rs b/crates/polars-core/src/datatypes/any_value.rs
@@ -367,6 +367,11 @@ impl AnyValue<'static> {
             _ => AnyValue::Null,
         }
     }
+
+    /// Can the [`AnyValue`] exist as having `dtype` as its `DataType`.
+    pub fn can_have_dtype(&self, dtype: &DataType) -> bool {
+        matches!(self, AnyValue::Null) || dtype == &self.dtype()
+    }
 }
 
 impl<'a> AnyValue<'a> {

diff --git a/crates/polars-core/src/frame/column/arithmetic.rs b/crates/polars-core/src/frame/column/arithmetic.rs
@@ -40,7 +40,6 @@ fn op_with_broadcast<F: Fn(&Series, &Series) -> PolarsResult<Series>>(
 
     let length = output_length(l, r)?;
     match (l, r) {
-        (Column::Series(l), Column::Series(r)) => op(l, r).map(Column::from),
         (Column::Series(l), Column::Scalar(r)) => {
             let r = r.as_single_value_series();
             if l.len() == 1 {
@@ -63,6 +62,7 @@ fn op_with_broadcast<F: Fn(&Series, &Series) -> PolarsResult<Series>>(
             op,
             length,
         ),
+        (l, r) => op(l.as_materialized_series(), r.as_materialized_series()).map(Column::from),
     }
 }
 
@@ -73,6 +73,8 @@ fn num_op_with_broadcast<T: Num + NumCast, F: Fn(&Series, T) -> Series>(
 ) -> Column {
     match c {
         Column::Series(s) => op(s, n).into(),
+        // @partition-opt
+        Column::Partitioned(s) => op(s.as_materialized_series(), n).into(),
         Column::Scalar(s) => {
             ScalarColumn::from_single_value_series(op(&s.as_single_value_series(), n), s.len())
                 .into()

diff --git a/crates/polars-core/src/frame/column/compare.rs b/crates/polars-core/src/frame/column/compare.rs
@@ -5,7 +5,6 @@ use super::{BooleanChunked, ChunkCompareEq, ChunkCompareIneq, ChunkExpandAtIndex
 macro_rules! column_element_wise_broadcasting {
     ($lhs:expr, $rhs:expr, $op:expr) => {
         match ($lhs, $rhs) {
-            (Column::Series(lhs), Column::Series(rhs)) => $op(lhs, rhs),
             (Column::Series(lhs), Column::Scalar(rhs)) => $op(lhs, &rhs.as_single_value_series()),
             (Column::Scalar(lhs), Column::Series(rhs)) => $op(&lhs.as_single_value_series(), rhs),
             (Column::Scalar(lhs), Column::Scalar(rhs)) => {
@@ -17,6 +16,7 @@ macro_rules! column_element_wise_broadcasting {
                     }
                 })
             },
+            (lhs, rhs) => $op(lhs.as_materialized_series(), rhs.as_materialized_series()),
         }
     };
 }