feat: IR serde

crates/polars-core/src/frame/explode.rs

@@ -1,5 +1,7 @@
 use arrow::offset::OffsetsBuffer;
 use rayon::prelude::*;
+#[cfg(feature = "serde")]
+use serde::{Deserialize, Serialize};
 use smartstring::alias::String as SmartString;
 
 use crate::chunked_array::ops::explode::offsets_to_indexes;
@@ -18,6 +20,7 @@ fn get_exploded(series: &Series) -> PolarsResult<(Series, OffsetsBuffer<i64>)> {
 
 /// Arguments for `[DataFrame::unpivot]` function
 #[derive(Clone, Default, Debug, PartialEq, Eq, Hash)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 pub struct UnpivotArgsIR {
     pub on: Vec<SmartString>,
     pub index: Vec<SmartString>,

crates/polars-plan/Cargo.toml

@@ -184,6 +184,7 @@ offset_by = ["polars-time/offset_by"]
 
 bigidx = ["polars-core/bigidx"]
 polars_cloud = ["serde", "ciborium"]
+ir_serde = ["serde", "polars-utils/ir_serde"]
 
 panic_on_schema = []
 

crates/polars-plan/src/plans/aexpr/mod.rs

@@ -11,6 +11,8 @@ use polars_core::chunked_array::cast::CastOptions;
 use polars_core::prelude::*;
 use polars_core::utils::{get_time_units, try_get_supertype};
 use polars_utils::arena::{Arena, Node};
+#[cfg(feature = "ir_serde")]
+use serde::{Deserialize, Serialize};
 use strum_macros::IntoStaticStr;
 pub use utils::*;
 
@@ -19,6 +21,7 @@ use crate::plans::Context;
 use crate::prelude::*;
 
 #[derive(Clone, Debug, IntoStaticStr)]
+#[cfg_attr(feature = "ir_serde", derive(Serialize, Deserialize))]
 pub enum IRAggExpr {
     Min {
         input: Node,
@@ -125,6 +128,7 @@ impl From<IRAggExpr> for GroupByMethod {
 
 /// IR expression node that is allocated in an [`Arena`][polars_utils::arena::Arena].
 #[derive(Clone, Debug, Default)]
+#[cfg_attr(feature = "ir_serde", derive(Serialize, Deserialize))]
 pub enum AExpr {
     Explode(Node),
     Alias(Node, ColumnName),
@@ -164,6 +168,7 @@ pub enum AExpr {
         truthy: Node,
         falsy: Node,
     },
+    #[cfg_attr(feature = "ir_serde", serde(skip))]
     AnonymousFunction {
         input: Vec<ExprIR>,
         function: SpecialEq<Arc<dyn SeriesUdf>>,

crates/polars-plan/src/plans/builder_dsl.rs

@@ -442,7 +442,7 @@ impl DslBuilder {
         function: F,
         optimizations: AllowedOptimizations,
         schema: Option<Arc<dyn UdfSchema>>,
-        name: &'static str,
+        name: &str,
     ) -> Self
     where
         F: DataFrameUdf + 'static,
@@ -457,7 +457,7 @@ impl DslBuilder {
                 predicate_pd: optimizations.contains(OptState::PREDICATE_PUSHDOWN),
                 projection_pd: optimizations.contains(OptState::PROJECTION_PUSHDOWN),
                 streamable: optimizations.contains(OptState::STREAMING),
-                fmt_str: name,
+                fmt_str: name.into(),
             }),
         }
         .into()

crates/polars-plan/src/plans/expr_ir.rs

@@ -3,10 +3,14 @@ use std::hash::Hash;
 #[cfg(feature = "cse")]
 use std::hash::Hasher;
 
+#[cfg(feature = "ir_serde")]
+use serde::{Deserialize, Serialize};
+
 use super::*;
 use crate::constants::{get_len_name, LITERAL_NAME};
 
 #[derive(Default, Debug, Clone, Hash, PartialEq, Eq)]
+#[cfg_attr(feature = "ir_serde", derive(Serialize, Deserialize))]
 pub enum OutputName {
     /// No not yet set.
     #[default]
@@ -40,6 +44,7 @@ impl OutputName {
 }
 
 #[derive(Clone, Debug, PartialEq, Eq)]
+#[cfg_attr(feature = "ir_serde", derive(Serialize, Deserialize))]
 pub struct ExprIR {
     /// Output name of this expression.
     output_name: OutputName,

crates/polars-plan/src/plans/functions/mod.rs

@@ -26,11 +26,13 @@ use crate::dsl::python_udf::PythonFunction;
 use crate::plans::functions::merge_sorted::merge_sorted;
 use crate::prelude::*;
 
+#[cfg_attr(feature = "ir_serde", derive(Serialize, Deserialize))]
 #[derive(Clone, IntoStaticStr)]
 #[strum(serialize_all = "SCREAMING_SNAKE_CASE")]
 pub enum FunctionIR {
     #[cfg(feature = "python")]
     OpaquePython(OpaquePythonUdf),
+    #[cfg_attr(feature = "ir_serde", serde(skip))]
     Opaque {
         function: Arc<dyn DataFrameUdf>,
         schema: Option<Arc<dyn UdfSchema>>,
@@ -40,14 +42,15 @@ pub enum FunctionIR {
         projection_pd: bool,
         streamable: bool,
         // used for formatting
-        fmt_str: &'static str,
+        fmt_str: String,
     },
     FastCount {
         paths: Arc<Vec<PathBuf>>,
         scan_type: FileScan,
         alias: Option<Arc<str>>,
     },
     /// Streaming engine pipeline
+    #[cfg_attr(feature = "ir_serde", serde(skip))]
     Pipeline {
         function: Arc<Mutex<dyn DataFrameUdfMut>>,
         schema: SchemaRef,
@@ -71,20 +74,24 @@ pub enum FunctionIR {
         new: Arc<[SmartString]>,
         // A column name gets swapped with an existing column
         swapping: bool,
+        #[cfg_attr(feature = "ir_serde", serde(skip))]
         schema: CachedSchema,
     },
     Explode {
         columns: Arc<[ColumnName]>,
+        #[cfg_attr(feature = "ir_serde", serde(skip))]
         schema: CachedSchema,
     },
     #[cfg(feature = "pivot")]
     Unpivot {
         args: Arc<UnpivotArgsIR>,
+        #[cfg_attr(feature = "ir_serde", serde(skip))]
         schema: CachedSchema,
     },
     RowIndex {
         name: Arc<str>,
         // Might be cached.
+        #[cfg_attr(feature = "ir_serde", serde(skip))]
         schema: CachedSchema,
         offset: Option<IdxSize>,
     },

crates/polars-plan/src/plans/ir/mod.rs

@@ -14,6 +14,8 @@ use hive::HivePartitions;
 use polars_core::prelude::*;
 use polars_utils::idx_vec::UnitVec;
 use polars_utils::unitvec;
+#[cfg(feature = "ir_serde")]
+use serde::{Deserialize, Serialize};
 
 use crate::prelude::*;
 
@@ -33,6 +35,7 @@ pub struct IRPlanRef<'a> {
 /// [`IR`] is a representation of [`DslPlan`] with [`Node`]s which are allocated in an [`Arena`]
 /// In this IR the logical plan has access to the full dataset.
 #[derive(Clone, Debug, Default)]
+#[cfg_attr(feature = "ir_serde", derive(Serialize, Deserialize))]
 pub enum IR {
     #[cfg(feature = "python")]
     PythonScan {
@@ -105,6 +108,7 @@ pub enum IR {
         keys: Vec<ExprIR>,
         aggs: Vec<ExprIR>,
         schema: SchemaRef,
+        #[cfg_attr(feature = "ir_serde", serde(skip))]
         apply: Option<Arc<dyn DataFrameUdf>>,
         maintain_order: bool,
         options: Arc<GroupbyOptions>,

crates/polars-plan/src/plans/options.rs

@@ -85,6 +85,7 @@ pub struct DistinctOptionsDSL {
 }
 
 #[derive(Clone, Debug, Eq, PartialEq, Hash)]
+#[cfg_attr(feature = "ir_serde", derive(Serialize, Deserialize))]
 pub struct DistinctOptionsIR {
     /// Subset of columns that will be taken into account.
     pub subset: Option<Arc<[ColumnName]>>,

crates/polars-utils/Cargo.toml

@@ -21,6 +21,7 @@ num-traits = { workspace = true }
 once_cell = { workspace = true }
 raw-cpuid = { workspace = true }
 rayon = { workspace = true }
+serde = { workspace = true, optional = true }
 smartstring = { workspace = true }
 stacker = { workspace = true }
 sysinfo = { version = "0.31", default-features = false, features = ["system"], optional = true }
@@ -35,3 +36,4 @@ version_check = { workspace = true }
 mmap = ["memmap"]
 bigidx = []
 nightly = []
+ir_serde = ["serde"]

crates/polars-utils/src/arena.rs

@@ -1,5 +1,8 @@
 use std::sync::atomic::{AtomicU32, Ordering};
 
+#[cfg(feature = "ir_serde")]
+use serde::{Deserialize, Serialize};
+
 use crate::error::*;
 use crate::slice::GetSaferUnchecked;
 
@@ -21,6 +24,7 @@ fn index_of<T>(slice: &[T], item: &T) -> Option<usize> {
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Ord, PartialOrd)]
 #[repr(transparent)]
+#[cfg_attr(feature = "ir_serde", derive(Serialize, Deserialize))]
 pub struct Node(pub usize);
 
 impl Default for Node {
@@ -32,6 +36,7 @@ impl Default for Node {
 static ARENA_VERSION: AtomicU32 = AtomicU32::new(0);
 
 #[derive(Debug, Clone)]
+#[cfg_attr(feature = "ir_serde", derive(Serialize, Deserialize))]
 pub struct Arena<T> {
     version: u32,
     items: Vec<T>,

crates/polars/Cargo.toml

@@ -228,6 +228,7 @@ zip_with = ["polars-core/zip_with"]
 
 bigidx = ["polars-core/bigidx", "polars-lazy?/bigidx", "polars-ops/big_idx"]
 polars_cloud = ["polars-lazy?/polars_cloud"]
+ir_serde = ["polars-plan/ir_serde"]
 
 test = [
   "lazy",