fix(python, rust): fix serde of expression (pickle) (#5333)

pola-rs · Oct 25, 2022 · 4a5a1b2 · 4a5a1b2
1 parent bf77257
commit 4a5a1b2
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 14 deletions.
diff --git a/polars/polars-lazy/polars-plan/src/dsl/expr.rs b/polars/polars-lazy/polars-plan/src/dsl/expr.rs
@@ -312,16 +312,6 @@ pub enum Expr {
         truthy: Box<Expr>,
         falsy: Box<Expr>,
     },
-    #[cfg_attr(feature = "serde", serde(skip))]
-    AnonymousFunction {
-        /// function arguments
-        input: Vec<Expr>,
-        /// function to apply
-        function: SpecialEq<Arc<dyn SeriesUdf>>,
-        /// output dtype of the function
-        output_type: GetOutput,
-        options: FunctionOptions,
-    },
     Function {
         /// function arguments
         input: Vec<Expr>,
@@ -353,15 +343,26 @@ pub enum Expr {
     Exclude(Box<Expr>, Vec<Excluded>),
     /// Set root name as Alias
     KeepName(Box<Expr>),
+    /// Special case that does not need columns
+    Count,
+    /// Take the nth column in the `DataFrame`
+    Nth(i64),
+    // skipped fields must be last otherwise serde fails in pickle
     #[cfg_attr(feature = "serde", serde(skip))]
     RenameAlias {
         function: SpecialEq<Arc<dyn RenameAliasFn>>,
         expr: Box<Expr>,
     },
-    /// Special case that does not need columns
-    Count,
-    /// Take the nth column in the `DataFrame`
-    Nth(i64),
+    #[cfg_attr(feature = "serde", serde(skip))]
+    AnonymousFunction {
+        /// function arguments
+        input: Vec<Expr>,
+        /// function to apply
+        function: SpecialEq<Arc<dyn SeriesUdf>>,
+        /// output dtype of the function
+        output_type: GetOutput,
+        options: FunctionOptions,
+    },
 }
 
 // TODO! derive. This is only a temporary fix

diff --git a/py-polars/tests/unit/io/test_pickle.py b/py-polars/tests/unit/io/test_pickle.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import io
 import pickle
 
 import polars as pl
@@ -14,3 +15,12 @@ def test_pickle() -> None:
     b = pickle.dumps(df)
     out = pickle.loads(b)
     assert df.frame_equal(out, null_equal=True)
+
+
+def test_pickle_expr() -> None:
+    for e in [pl.all(), pl.count()]:
+        f = io.BytesIO()
+        pickle.dump(e, f)
+
+        f.seek(0)
+        pickle.load(f)