python: fix to_list wo/ pyarrow

pola-rs · Nov 5, 2021 · 796fb7d · 796fb7d
1 parent 14c9c1d
commit 796fb7d
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 16 deletions.
diff --git a/py-polars/polars/eager/series.py b/py-polars/polars/eager/series.py
@@ -1694,7 +1694,8 @@ def to_list(self) -> tp.List[Optional[Any]]:
         <class 'list'>
 
         """
-        if self.dtype != Object:
+        # maybe we should not use pyarrow at all.
+        if (self.dtype != Object) and _PYARROW_AVAILABLE:
             return self.to_arrow().to_pylist()
         return self._s.to_list()
 

diff --git a/py-polars/src/series.rs b/py-polars/src/series.rs
@@ -630,21 +630,26 @@ impl PySeries {
 
         let series = &self.series;
 
+        let primitive_to_list = |dt: &DataType, series: &Series| {
+            match dt {
+                DataType::Boolean => PyList::new(python, series.bool().unwrap()),
+                DataType::Utf8 => PyList::new(python, series.utf8().unwrap()),
+                DataType::UInt8 => PyList::new(python, series.u8().unwrap()),
+                DataType::UInt16 => PyList::new(python, series.u16().unwrap()),
+                DataType::UInt32 => PyList::new(python, series.u32().unwrap()),
+                DataType::UInt64 => PyList::new(python, series.u64().unwrap()),
+                DataType::Int8 => PyList::new(python, series.i8().unwrap()),
+                DataType::Int16 => PyList::new(python, series.i16().unwrap()),
+                DataType::Int32 => PyList::new(python, series.i32().unwrap()),
+                DataType::Int64 => PyList::new(python, series.i64().unwrap()),
+                DataType::Float32 => PyList::new(python, series.f32().unwrap()),
+                DataType::Float64 => PyList::new(python, series.f64().unwrap()),
+                dt => panic!("to_list() not implemented for {:?}", dt)
+            }
+        };
+
         let pylist = match series.dtype() {
-            DataType::Boolean => PyList::new(python, series.bool().unwrap()),
-            DataType::Utf8 => PyList::new(python, series.utf8().unwrap()),
-            DataType::UInt8 => PyList::new(python, series.u8().unwrap()),
-            DataType::UInt16 => PyList::new(python, series.u16().unwrap()),
-            DataType::UInt32 => PyList::new(python, series.u32().unwrap()),
-            DataType::UInt64 => PyList::new(python, series.u64().unwrap()),
-            DataType::Int8 => PyList::new(python, series.i8().unwrap()),
-            DataType::Int16 => PyList::new(python, series.i16().unwrap()),
-            DataType::Int32 => PyList::new(python, series.i32().unwrap()),
-            DataType::Int64 => PyList::new(python, series.i64().unwrap()),
-            DataType::Float32 => PyList::new(python, series.f32().unwrap()),
-            DataType::Float64 => PyList::new(python, series.f64().unwrap()),
-            DataType::Date => PyList::new(python, &series.date().unwrap().0),
-            DataType::Datetime => PyList::new(python, &series.datetime().unwrap().0),
+            DataType::Categorical => PyList::new(python, series.categorical().unwrap()),
             DataType::Object(_) => {
                 let v = PyList::empty(python);
                 for i in 0..series.len() {
@@ -655,7 +660,25 @@ impl PySeries {
                 }
                 v
             }
-            dt => panic!("to_list() not implemented for {:?}", dt),
+            DataType::List(inner_dtype) => {
+                let v = PyList::empty(python);
+                let ca = series.list().unwrap();
+                for opt_s in ca.amortized_iter() {
+                    match opt_s {
+                        None => {
+                            v.append(python.None()).unwrap();
+                        }
+                        Some(s) => {
+                            let pylst = primitive_to_list(&inner_dtype, s.as_ref());
+                            v.append(pylst).unwrap();
+                        }
+                    }
+                }
+                v
+            }
+            DataType::Date => PyList::new(python, &series.date().unwrap().0),
+            DataType::Datetime => PyList::new(python, &series.datetime().unwrap().0),
+            dt => primitive_to_list(dt, series)
         };
         pylist.to_object(python)
     }