python take collect into alignedvec

pola-rs · Aug 5, 2021 · f733ea9 · f733ea9
1 parent 32a71e7
commit f733ea9
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 5 deletions.
diff --git a/polars/polars-core/src/chunked_array/mod.rs b/polars/polars-core/src/chunked_array/mod.rs
@@ -35,7 +35,6 @@ pub mod strings;
 pub mod temporal;
 mod trusted_len;
 pub mod upstream_traits;
-
 use arrow::array::Array;
 
 use crate::chunked_array::builder::categorical::RevMapping;
@@ -868,6 +867,12 @@ pub(crate) fn to_array<T: PolarsPrimitiveType>(
     Arc::new(to_primitive::<T>(values, validity))
 }
 
+impl<T: PolarsNumericType> From<PrimitiveArray<T::Native>> for ChunkedArray<T> {
+    fn from(a: PrimitiveArray<T::Native>) -> Self {
+        ChunkedArray::new_from_chunks("", vec![Arc::new(a)])
+    }
+}
+
 #[cfg(test)]
 pub(crate) mod test {
     use crate::prelude::*;

diff --git a/py-polars/src/conversion.rs b/py-polars/src/conversion.rs
@@ -11,6 +11,7 @@ use pyo3::types::PySequence;
 use pyo3::{PyAny, PyResult};
 use std::fmt::{Display, Formatter};
 use std::hash::{Hash, Hasher};
+use polars_core::utils::arrow::datatypes::ArrowNativeType;
 
 #[repr(transparent)]
 pub struct Wrap<T>(pub T);
@@ -268,3 +269,14 @@ impl Default for ObjectValue {
         }
     }
 }
+
+impl<'a, T: ArrowNativeType + FromPyObject<'a>> FromPyObject<'a> for Wrap<AlignedVec<T>> {
+    fn extract(obj: &'a PyAny) -> PyResult<Self> {
+        let seq = <PySequence as PyTryFrom>::try_from(obj)?;
+        let mut v = AlignedVec::with_capacity(seq.len().unwrap_or(0) as usize);
+        for item in seq.iter()? {
+            v.push(item?.extract::<T>()?);
+        }
+        Ok(Wrap(v))
+    }
+}
diff --git a/py-polars/src/dataframe.rs b/py-polars/src/dataframe.rs
@@ -514,10 +514,12 @@ impl PyDataFrame {
         }
     }
 
-    pub fn take(&self, indices: Vec<u32>) -> PyResult<Self> {
+    pub fn take(&self, indices: Wrap<AlignedVec<u32>>) -> PyResult<Self> {
+        let indices = indices.0;
+        let indices = indices.into_primitive_array::<UInt32Type>(None);
         let df = self
             .df
-            .take_iter(indices.iter().map(|i| *i as usize))
+            .take(&indices.into())
             .map_err(PyPolarsEr::from)?;
         Ok(PyDataFrame::new(df))
     }

diff --git a/py-polars/src/series.rs b/py-polars/src/series.rs
@@ -388,10 +388,13 @@ impl PySeries {
         self.series.arg_max()
     }
 
-    pub fn take(&self, indices: Vec<usize>) -> PyResult<Self> {
+    pub fn take(&self, indices: Wrap<AlignedVec<u32>>) -> PyResult<Self> {
+        let indices = indices.0;
+        let indices = indices.into_primitive_array::<UInt32Type>(None);
+
         let take = self
             .series
-            .take_iter(&mut indices.iter().copied())
+            .take(&indices.into())
             .map_err(PyPolarsEr::from)?;
         Ok(PySeries::new(take))
     }