Dev

polars/src/chunked_array/object/mod.rs

@@ -142,11 +142,9 @@ where
     T: Any + Debug + Clone + Send + Sync + Default,
 {
     pub fn get_as_any(&self, index: usize) -> &dyn Any {
+        let chunks = self.downcast_chunks();
         let (chunk_idx, idx) = self.index_to_chunked_index(index);
-        let arr = unsafe {
-            let arr = &**self.chunks.get_unchecked(chunk_idx);
-            &*(arr as *const dyn Array as *const ObjectArray<T>)
-        };
+        let arr = unsafe { *chunks.get_unchecked(chunk_idx) };
         arr.value(idx)
     }
 }

polars/src/chunked_array/ops/apply.rs

@@ -3,6 +3,30 @@ use crate::prelude::*;
 use crate::utils::Xob;
 use arrow::array::{ArrayRef, PrimitiveArray, StringArray};
 
+macro_rules! apply {
+    ($self:expr, $f:expr) => {{
+        if $self.null_count() == 0 {
+            $self.into_no_null_iter().map($f).collect()
+        } else {
+            $self.into_iter().map(|opt_v| opt_v.map($f)).collect()
+        }
+    }};
+}
+
+macro_rules! apply_enumerate {
+    ($self:expr, $f:expr) => {{
+        if $self.null_count() == 0 {
+            $self.into_no_null_iter().enumerate().map($f).collect()
+        } else {
+            $self
+                .into_iter()
+                .enumerate()
+                .map(|(idx, opt_v)| opt_v.map(|v| $f((idx, v))))
+                .collect()
+        }
+    }};
+}
+
 impl<'a, T> ChunkApply<'a, T::Native, T::Native> for ChunkedArray<T>
 where
     T: PolarsNumericType,
@@ -40,31 +64,71 @@ where
             ca
         }
     }
+
+    fn apply_with_idx<F>(&'a self, f: F) -> Self
+    where
+        F: Fn((usize, T::Native)) -> T::Native + Copy,
+    {
+        if self.null_count() == 0 {
+            let ca: Xob<_> = self.into_no_null_iter().enumerate().map(f).collect();
+            ca.into_inner()
+        } else {
+            self.into_iter()
+                .enumerate()
+                .map(|(idx, opt_v)| opt_v.map(|v| f((idx, v))))
+                .collect()
+        }
+    }
+
+    fn apply_with_idx_on_opt<F>(&'a self, f: F) -> Self
+    where
+        F: Fn((usize, Option<T::Native>)) -> Option<T::Native> + Copy,
+    {
+        self.into_iter().enumerate().map(f).collect()
+    }
 }
 
 impl<'a> ChunkApply<'a, bool, bool> for BooleanChunked {
     fn apply<F>(&self, f: F) -> Self
     where
         F: Fn(bool) -> bool + Copy,
     {
-        if self.null_count() == 0 {
-            self.into_no_null_iter().map(f).collect()
-        } else {
-            self.into_iter().map(|opt_v| opt_v.map(|v| f(v))).collect()
-        }
+        apply!(self, f)
+    }
+    fn apply_with_idx<F>(&'a self, f: F) -> Self
+    where
+        F: Fn((usize, bool)) -> bool + Copy,
+    {
+        apply_enumerate!(self, f)
+    }
+
+    fn apply_with_idx_on_opt<F>(&'a self, f: F) -> Self
+    where
+        F: Fn((usize, Option<bool>)) -> Option<bool> + Copy,
+    {
+        self.into_iter().enumerate().map(f).collect()
     }
 }
 
 impl<'a> ChunkApply<'a, &'a str, String> for Utf8Chunked {
     fn apply<F>(&'a self, f: F) -> Self
     where
-        F: Fn(&'a str) -> String,
+        F: Fn(&'a str) -> String + Copy,
     {
-        if self.null_count() == 0 {
-            self.into_no_null_iter().map(f).collect()
-        } else {
-            self.into_iter().map(|opt_v| opt_v.map(|v| f(v))).collect()
-        }
+        apply!(self, f)
+    }
+    fn apply_with_idx<F>(&'a self, f: F) -> Self
+    where
+        F: Fn((usize, &'a str)) -> String + Copy,
+    {
+        apply_enumerate!(self, f)
+    }
+
+    fn apply_with_idx_on_opt<F>(&'a self, f: F) -> Self
+    where
+        F: Fn((usize, Option<&'a str>)) -> Option<String> + Copy,
+    {
+        self.into_iter().enumerate().map(f).collect()
     }
 }
 

polars/src/chunked_array/ops/chunkops.rs

@@ -1,7 +1,10 @@
 use crate::chunked_array::builder::get_list_builder;
+use crate::chunked_array::object::builder::ObjectChunkedBuilder;
 use crate::prelude::*;
 use arrow::array::{Array, ArrayRef, PrimitiveBuilder, StringBuilder};
 use arrow::compute::concat;
+use std::any::Any;
+use std::fmt::Debug;
 use std::sync::Arc;
 
 pub trait ChunkOps {
@@ -189,12 +192,46 @@ impl ChunkOps for ListChunked {
     }
 }
 
-impl<T> ChunkOps for ObjectChunked<T> {
-    fn rechunk(&self, _chunk_lengths: Option<&[usize]>) -> Result<Self>
+impl<T> ChunkOps for ObjectChunked<T>
+where
+    T: Any + Debug + Clone + Send + Sync + Default,
+{
+    fn rechunk(&self, chunk_lengths: Option<&[usize]>) -> Result<Self>
     where
         Self: std::marker::Sized,
     {
-        todo!()
+        match (self.chunks.len(), chunk_lengths.map(|v| v.len())) {
+            (1, Some(1)) | (1, None) => Ok(self.clone()),
+            (1, Some(_)) => mimic_chunks(&self.chunks[0], chunk_lengths.unwrap(), self.name()),
+            (_, None) => {
+                let mut builder = ObjectChunkedBuilder::new(self.name(), self.len());
+                let chunks = self.downcast_chunks();
+
+                // todo! use iterators once implemented
+                // no_null path
+                if self.null_count() == 0 {
+                    for idx in 0..self.len() {
+                        let (chunk_idx, idx) = self.index_to_chunked_index(idx);
+                        let arr = unsafe { &**chunks.get_unchecked(chunk_idx) };
+                        builder.append_value(arr.value(idx).clone())
+                    }
+                } else {
+                    for idx in 0..self.len() {
+                        let (chunk_idx, idx) = self.index_to_chunked_index(idx);
+                        let arr = unsafe { &**chunks.get_unchecked(chunk_idx) };
+                        if arr.is_valid(idx) {
+                            builder.append_value(arr.value(idx).clone())
+                        } else {
+                            builder.append_null()
+                        }
+                    }
+                }
+                Ok(builder.finish())
+            }
+            _ => Err(PolarsError::Other(
+                "rechunk of ObjectChunked still needs to be implemented".into(),
+            )),
+        }
     }
 
     fn optional_rechunk<A>(&self, _rhs: &ChunkedArray<A>) -> Result<Option<Self>>

polars/src/chunked_array/ops/mod.rs

@@ -332,6 +332,16 @@ pub trait ChunkApply<'a, A, B> {
     fn apply<F>(&'a self, f: F) -> Self
     where
         F: Fn(A) -> B + Copy;
+
+    /// Apply a closure elementwise. The closure gets the index of the element as first argument.
+    fn apply_with_idx<F>(&'a self, f: F) -> Self
+    where
+        F: Fn((usize, A)) -> B + Copy;
+
+    /// Apply a closure elementwise. The closure gets the index of the element as first argument.
+    fn apply_with_idx_on_opt<F>(&'a self, f: F) -> Self
+    where
+        F: Fn((usize, Option<A>)) -> Option<B> + Copy;
 }
 
 /// Aggregation operations