feat: Implement serde for Null columns (#17218)

pola-rs · Jun 26, 2024 · 053765b · 053765b
1 parent 3115865
commit 053765b
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 4 deletions.
diff --git a/crates/polars-core/src/serde/chunked_array.rs b/crates/polars-core/src/serde/chunked_array.rs
@@ -5,6 +5,7 @@ use serde::{Serialize, Serializer};
 
 use crate::chunked_array::metadata::MetadataFlags;
 use crate::prelude::*;
+use crate::series::implementations::null::NullChunked;
 
 pub struct IterSer<I>
 where
@@ -174,3 +175,21 @@ impl Serialize for StructChunked {
         }
     }
 }
+
+impl Serialize for NullChunked {
+    fn serialize<S>(
+        &self,
+        serializer: S,
+    ) -> std::result::Result<<S as Serializer>::Ok, <S as Serializer>::Error>
+    where
+        S: Serializer,
+    {
+        {
+            let mut state = serializer.serialize_map(Some(3))?;
+            state.serialize_entry("name", self.name())?;
+            state.serialize_entry("datatype", self.dtype())?;
+            state.serialize_entry("values", &IterSer::new(std::iter::once(self.len())))?;
+            state.end()
+        }
+    }
+}
diff --git a/crates/polars-core/src/serde/series.rs b/crates/polars-core/src/serde/series.rs
@@ -75,6 +75,10 @@ impl Serialize for Series {
                 let ca = self.decimal().unwrap();
                 ca.serialize(serializer)
             },
+            DataType::Null => {
+                let ca = self.null().unwrap();
+                ca.serialize(serializer)
+            },
             dt => {
                 with_match_physical_numeric_polars_type!(dt, |$T| {
                 let ca: &ChunkedArray<$T> = self.as_ref().as_ref().as_ref();
@@ -276,6 +280,11 @@ impl<'de> Deserialize<'de> for Series {
                         let values: Vec<Option<Cow<str>>> = map.next_value()?;
                         Ok(Series::new(&name, values).cast(&dt).unwrap())
                     },
+                    DataType::Null => {
+                        let values: Vec<usize> = map.next_value()?;
+                        let len = values.first().unwrap();
+                        Ok(Series::new_null(&name, *len))
+                    },
                     dt => {
                         panic!("{dt:?} dtype deserialization not yet implemented")
                     },

diff --git a/py-polars/tests/unit/dataframe/test_serde.py b/py-polars/tests/unit/dataframe/test_serde.py
@@ -20,7 +20,6 @@
 @given(
     df=dataframes(
         excluded_dtypes=[
-            pl.Null,  # Not implemented yet
             pl.Float32,  # Bug, see: https://github.com/pola-rs/polars/issues/17211
             pl.Float64,  # Bug, see: https://github.com/pola-rs/polars/issues/17211
         ],
@@ -170,9 +169,8 @@ def test_df_serde_float_inf_nan() -> None:
     assert_frame_equal(result, df)
 
 
-@pytest.mark.xfail(reason="Not implemented yet")
 def test_df_serde_null() -> None:
-    df = pl.DataFrame({"a": [None, None]})
+    df = pl.DataFrame({"a": [None, None]}, schema={"a": pl.Null})
     ser = df.serialize()
     result = pl.DataFrame.deserialize(io.StringIO(ser))
     assert_frame_equal(result, df)

diff --git a/py-polars/tests/unit/lazyframe/test_serde.py b/py-polars/tests/unit/lazyframe/test_serde.py
@@ -18,7 +18,6 @@
     lf=dataframes(
         lazy=True,
         excluded_dtypes=[
-            pl.Null,  # Not implemented yet
             pl.Float32,  # Bug, see: https://github.com/pola-rs/polars/issues/17211
             pl.Float64,  # Bug, see: https://github.com/pola-rs/polars/issues/17211
         ],