fix invalid dtype when taking from empty list

pola-rs · Nov 19, 2021 · 2c6f811 · 2c6f811
1 parent 28d4411
commit 2c6f811
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 23 deletions.
diff --git a/polars/polars-core/src/chunked_array/ops/full.rs b/polars/polars-core/src/chunked_array/ops/full.rs
@@ -23,9 +23,8 @@ where
     T: PolarsNumericType,
 {
     fn full_null(name: &str, length: usize) -> Self {
-        let mut ca = (0..length).map(|_| None).collect::<Self>();
-        ca.rename(name);
-        ca
+        let arr = new_null_array(T::get_dtype().to_arrow(), length).into();
+        ChunkedArray::new_from_chunks(name, vec![arr])
     }
 }
 impl ChunkFull<bool> for BooleanChunked {
@@ -38,9 +37,8 @@ impl ChunkFull<bool> for BooleanChunked {
 
 impl ChunkFullNull for BooleanChunked {
     fn full_null(name: &str, length: usize) -> Self {
-        let mut ca = (0..length).map(|_| None).collect::<Self>();
-        ca.rename(name);
-        ca
+        let arr = new_null_array(DataType::Boolean.to_arrow(), length).into();
+        BooleanChunked::new_from_chunks(name, vec![arr])
     }
 }
 
@@ -57,11 +55,8 @@ impl<'a> ChunkFull<&'a str> for Utf8Chunked {
 
 impl ChunkFullNull for Utf8Chunked {
     fn full_null(name: &str, length: usize) -> Self {
-        let mut ca = (0..length)
-            .map::<Option<String>, _>(|_| None)
-            .collect::<Self>();
-        ca.rename(name);
-        ca
+        let arr = new_null_array(DataType::Utf8.to_arrow(), length).into();
+        Utf8Chunked::new_from_chunks(name, vec![arr])
     }
 }
 
@@ -77,11 +72,7 @@ impl ChunkFull<&Series> for ListChunked {
 
 impl ChunkFullNull for ListChunked {
     fn full_null(name: &str, length: usize) -> ListChunked {
-        let mut builder = ListBooleanChunkedBuilder::new(name, length, 0);
-        for _ in 0..length {
-            builder.append_null();
-        }
-        builder.finish()
+        ListChunked::full_null_with_dtype(name, length, &DataType::Boolean)
     }
 }
 
@@ -97,12 +88,21 @@ impl ChunkFullNull for CategoricalChunked {
 }
 
 impl ListChunked {
-    pub(crate) fn full_null_with_dtype(name: &str, length: usize, dt: &DataType) -> ListChunked {
-        let mut builder = get_list_builder(dt, 0, length, name);
-        for _ in 0..length {
-            builder.append_null();
-        }
-        builder.finish()
+    pub(crate) fn full_null_with_dtype(
+        name: &str,
+        length: usize,
+        inner_dtype: &DataType,
+    ) -> ListChunked {
+        let arr = new_null_array(
+            ArrowDataType::List(Box::new(ArrowField::new(
+                "item",
+                inner_dtype.to_arrow(),
+                true,
+            ))),
+            length,
+        )
+        .into();
+        ListChunked::new_from_chunks(name, vec![arr])
     }
 }
 

diff --git a/polars/polars-core/src/chunked_array/ops/take/mod.rs b/polars/polars-core/src/chunked_array/ops/take/mod.rs
@@ -343,7 +343,11 @@ impl ChunkTake for ListChunked {
         match indices {
             TakeIdx::Array(array) => {
                 if array.null_count() == array.len() {
-                    return Self::full_null(self.name(), array.len());
+                    return Self::full_null_with_dtype(
+                        self.name(),
+                        array.len(),
+                        &self.inner_dtype(),
+                    );
                 }
                 let array = match ca_self.chunks.len() {
                     1 => Arc::new(take_list_unchecked(chunks.next().unwrap(), array)) as ArrayRef,