diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/with_column.rs b/crates/nu_plugin_polars/src/dataframe/command/data/with_column.rs index 38bafc9f72..61f90be19f 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/with_column.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/with_column.rs @@ -104,9 +104,44 @@ impl PluginCommand for WithColumn { ), }, Example { - description: "Add series to the dataframe using a record", + description: "Add series to a lazyframe using a record", example: r#"[[a b]; [1 2] [3 4]] | polars into-lazy + | polars with-column { + c: ((polars col a) * 2) + d: ((polars col a) * 3) + } + | polars collect"#, + result: Some( + NuDataFrame::try_from_columns( + vec![ + Column::new( + "a".to_string(), + vec![Value::test_int(1), Value::test_int(3)], + ), + Column::new( + "b".to_string(), + vec![Value::test_int(2), Value::test_int(4)], + ), + Column::new( + "c".to_string(), + vec![Value::test_int(2), Value::test_int(6)], + ), + Column::new( + "d".to_string(), + vec![Value::test_int(3), Value::test_int(9)], + ), + ], + None, + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + Example { + description: "Add series to a dataframe using a record", + example: r#"[[a b]; [1 2] [3 4]] + | polars into-df | polars with-column { c: ((polars col a) * 2) d: ((polars col a) * 3) @@ -177,7 +212,7 @@ fn command_eager( if NuExpression::can_downcast(&new_column) { if let Some(name) = call.get_flag::>("name")? { return Err(ShellError::GenericError { - error: "Flag 'name' is unsuppored when used with expressions. Please use the `polars as` expression to name a column".into(), + error: "Flag 'name' is unsupported when used with expressions. Please use the `polars as` expression to name a column".into(), msg: "".into(), span: Some(name.span), help: Some("Use a `polars as` expression to name a column".into()), @@ -225,7 +260,7 @@ fn command_lazy( ) -> Result { if let Some(name) = call.get_flag::>("name")? { return Err(ShellError::GenericError { - error: "Flag 'name' is unsuppored for lazy dataframes. Please use the `polars as` expression to name a column".into(), + error: "Flag 'name' is unsupported for lazy dataframes. Please use the `polars as` expression to name a column".into(), msg: "".into(), span: Some(name.span), help: Some("Use a `polars as` expression to name a column".into()), diff --git a/crates/nu_plugin_polars/src/dataframe/values/nu_expression/mod.rs b/crates/nu_plugin_polars/src/dataframe/values/nu_expression/mod.rs index 60474139a6..b84f175bef 100644 --- a/crates/nu_plugin_polars/src/dataframe/values/nu_expression/mod.rs +++ b/crates/nu_plugin_polars/src/dataframe/values/nu_expression/mod.rs @@ -572,6 +572,7 @@ impl CustomValueSupport for NuExpression { match value { Value::Custom { val, .. } => val.as_any().downcast_ref::().is_some(), Value::List { vals, .. } => vals.iter().all(Self::can_downcast), + Value::Record { val, .. } => val.iter().all(|(_, value)| Self::can_downcast(value)), Value::String { .. } | Value::Int { .. } | Value::Bool { .. } | Value::Float { .. } => { true }