fix python

pola-rs · Aug 5, 2021 · c032df8 · c032df8
1 parent 965ff63
commit c032df8
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 9 deletions.
diff --git a/py-polars/src/conversion.rs b/py-polars/src/conversion.rs
@@ -4,7 +4,7 @@ use crate::series::PySeries;
 use polars::chunked_array::object::PolarsObjectSafe;
 use polars::frame::row::Row;
 use polars::prelude::AnyValue;
-use polars_core::utils::arrow::datatypes::ArrowNativeType;
+use polars_core::utils::arrow::types::NativeType;
 use pyo3::basic::CompareOp;
 use pyo3::conversion::{FromPyObject, IntoPy};
 use pyo3::prelude::*;
@@ -270,7 +270,7 @@ impl Default for ObjectValue {
     }
 }
 
-impl<'a, T: ArrowNativeType + FromPyObject<'a>> FromPyObject<'a> for Wrap<AlignedVec<T>> {
+impl<'a, T: NativeType + 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);

diff --git a/py-polars/src/dataframe.rs b/py-polars/src/dataframe.rs
@@ -516,8 +516,8 @@ impl PyDataFrame {
 
     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(&indices.into()).map_err(PyPolarsEr::from)?;
+        let indices = UInt32Chunked::new_from_aligned_vec("", indices);
+        let df = self.df.take(&indices).map_err(PyPolarsEr::from)?;
         Ok(PyDataFrame::new(df))
     }
 

diff --git a/py-polars/src/series.rs b/py-polars/src/series.rs
@@ -390,12 +390,9 @@ impl PySeries {
 
     pub fn take(&self, indices: Wrap<AlignedVec<u32>>) -> PyResult<Self> {
         let indices = indices.0;
-        let indices = indices.into_primitive_array::<UInt32Type>(None);
+        let indices = UInt32Chunked::new_from_aligned_vec("", indices);
 
-        let take = self
-            .series
-            .take(&indices.into())
-            .map_err(PyPolarsEr::from)?;
+        let take = self.series.take(&indices).map_err(PyPolarsEr::from)?;
         Ok(PySeries::new(take))
     }