doctest and rename schema

crates/polars-plan/src/logical_plan/functions/dsl.rs

@@ -90,6 +90,7 @@ impl DslFunction {
                     existing,
                     new,
                     swapping,
+                    schema: Default::default(),
                 }
             },
             DslFunction::Stats(_)

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

@@ -84,6 +84,8 @@ pub enum FunctionNode {
         new: Arc<[SmartString]>,
         // A column name gets swapped with an existing column
         swapping: bool,
+        #[cfg_attr(feature = "serde", serde(skip))]
+        schema: CachedSchema,
     },
     Explode {
         columns: Arc<[Arc<str>]>,
@@ -161,6 +163,7 @@ impl Hash for FunctionNode {
                 existing,
                 new,
                 swapping: _,
+                ..
             } => {
                 existing.hash(state);
                 new.hash(state);

crates/polars-plan/src/logical_plan/functions/schema.rs

@@ -8,7 +8,10 @@ impl FunctionNode {
         // We will likely add more branches later
         #[allow(clippy::single_match)]
         match self {
-            RowIndex { schema, .. } => {
+            RowIndex { schema, .. }
+            | Explode { schema, .. }
+            | Rename { schema, .. }
+            | Melt { schema, .. } => {
                 let mut guard = schema.lock().unwrap();
                 *guard = None;
             },
@@ -86,7 +89,12 @@ impl FunctionNode {
             },
             #[cfg(feature = "merge_sorted")]
             MergeSorted { .. } => Ok(Cow::Borrowed(input_schema)),
-            Rename { existing, new, .. } => rename_schema(input_schema, existing, new),
+            Rename {
+                existing,
+                new,
+                schema,
+                ..
+            } => rename_schema(input_schema, existing, new, schema),
             RowIndex { schema, name, .. } => {
                 Ok(Cow::Owned(row_index_schema(schema, input_schema, name)))
             },
@@ -191,7 +199,12 @@ fn rename_schema<'a>(
     input_schema: &'a SchemaRef,
     existing: &[SmartString],
     new: &[SmartString],
+    cached_schema: &CachedSchema,
 ) -> PolarsResult<Cow<'a, SchemaRef>> {
+    let mut guard = cached_schema.lock().unwrap();
+    if let Some(schema) = &*guard {
+        return Ok(Cow::Owned(schema.clone()));
+    }
     let mut new_schema = input_schema.iter_fields().collect::<Vec<_>>();
 
     for (old, new) in existing.iter().zip(new.iter()) {
@@ -201,5 +214,7 @@ fn rename_schema<'a>(
             new_schema[idx].name = new.as_str().into();
         }
     }
-    Ok(Cow::Owned(Arc::new(new_schema.into_iter().collect())))
+    let schema: SchemaRef = Arc::new(new_schema.into_iter().collect());
+    *guard = Some(schema.clone());
+    Ok(Cow::Owned(schema))
 }

crates/polars-plan/src/logical_plan/optimizer/projection_pushdown/functions/mod.rs

@@ -21,6 +21,7 @@ pub(super) fn process_functions(
             existing,
             new,
             swapping,
+            schema: _,
         } => {
             process_rename(
                 &mut acc_projections,

py-polars/polars/lazyframe/frame.py

@@ -582,7 +582,7 @@ def serialize(self, file: IOBase | str | Path | None = None) -> str | None:
         >>> lf = pl.LazyFrame({"a": [1, 2, 3]}).sum()
         >>> json = lf.serialize()
         >>> json
-        '{"Select":{"expr":[{"Agg":{"Sum":{"Column":"a"}}}],"input":{"DataFrameScan":{"df":{"columns":[{"name":"a","datatype":"Int64","bit_settings":"","values":[1,2,3]}]},"schema":{"inner":{"a":"Int64"}},"output_schema":null,"projection":null,"selection":null}},"schema":{"inner":{"a":"Int64"}},"options":{"run_parallel":true,"duplicate_check":true}}}'
+        '{"MapFunction":{"input":{"DataFrameScan":{"df":{"columns":[{"name":"a","datatype":"Int64","bit_settings":"","values":[1,2,3]}]},"schema":{"inner":{"a":"Int64"}},"output_schema":null,"projection":null,"selection":null}},"function":{"Stats":"Sum"}}}'
 
         The logical plan can later be deserialized back into a LazyFrame.