diff --git a/crates/nu_plugin_polars/src/dataframe/command/data/shift.rs b/crates/nu_plugin_polars/src/dataframe/command/data/shift.rs index bb58283b7e..ba2f5b4532 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/data/shift.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/data/shift.rs @@ -37,10 +37,16 @@ impl PluginCommand for Shift { "Expression used to fill the null values (lazy df)", Some('f'), ) - .input_output_type( - Type::Custom("dataframe".into()), - Type::Custom("dataframe".into()), - ) + .input_output_types(vec![ + ( + Type::Custom("dataframe".into()), + Type::Custom("dataframe".into()), + ), + ( + Type::Custom("expression".into()), + Type::Custom("expression".into()), + ), + ]) .category(Category::Custom("dataframe or lazyframe".into())) } @@ -82,6 +88,42 @@ impl PluginCommand for Shift { .into_value(Span::test_data()), ), }, + Example { + description: "Shift values of a column, fill absent values with 0", + example: "[[a]; [1] [2] [2] [3] [3]] + | polars into-lazy + | polars with-column {b: (polars col a | polars shift 2 --fill 0)} + | polars collect", + result: Some( + NuDataFrame::try_from_columns( + vec![ + Column::new( + "a".to_string(), + vec![ + Value::test_int(1), + Value::test_int(2), + Value::test_int(2), + Value::test_int(3), + Value::test_int(3), + ], + ), + Column::new( + "b".to_string(), + vec![ + Value::test_int(0), + Value::test_int(0), + Value::test_int(1), + Value::test_int(2), + Value::test_int(2), + ], + ), + ], + None, + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, ] } @@ -98,11 +140,28 @@ impl PluginCommand for Shift { match PolarsPluginObject::try_from_value(plugin, &value)? { PolarsPluginObject::NuDataFrame(df) => command_eager(plugin, engine, call, df), PolarsPluginObject::NuLazyFrame(lazy) => command_lazy(plugin, engine, call, lazy), + PolarsPluginObject::NuExpression(expr) => { + let shift: i64 = call.req(0)?; + let fill: Option = call.get_flag("fill")?; + + let res: NuExpression = match fill { + Some(ref fill) => { + let fill_expr = NuExpression::try_from_value(plugin, fill)?.into_polars(); + expr.into_polars() + .shift_and_fill(lit(shift), fill_expr) + .into() + } + None => expr.into_polars().shift(lit(shift)).into(), + }; + + res.to_pipeline_data(plugin, engine, call.head) + } _ => Err(cant_convert_err( &value, &[ PolarsPluginType::NuDataFrame, PolarsPluginType::NuLazyGroupBy, + PolarsPluginType::NuExpression, ], )), }