From eb0de25d190dc596cfbf055be9806741bc047c1b Mon Sep 17 00:00:00 2001
From: Jack Wright <56345+ayax79@users.noreply.github.com>
Date: Wed, 4 Sep 2024 03:19:29 -0700
Subject: [PATCH] Expression support for `polars strftime` (#13767)
# Description
Allows `polars strftime` to be used as an expression:
# User-Facing Changes
- `polars strftime` can now be used as an expression.
---
crates/nu_plugin_polars/Cargo.toml | 1 +
.../src/dataframe/series/string/strftime.rs | 111 +++++++++++++-----
2 files changed, 85 insertions(+), 27 deletions(-)
diff --git a/crates/nu_plugin_polars/Cargo.toml b/crates/nu_plugin_polars/Cargo.toml
index 2d3e1ee312..0c75a1fc56 100644
--- a/crates/nu_plugin_polars/Cargo.toml
+++ b/crates/nu_plugin_polars/Cargo.toml
@@ -73,6 +73,7 @@ features = [
"strings",
"string_to_integer",
"streaming",
+ "temporal",
"to_dummies",
]
optional = false
diff --git a/crates/nu_plugin_polars/src/dataframe/series/string/strftime.rs b/crates/nu_plugin_polars/src/dataframe/series/string/strftime.rs
index 85d3aa84a0..ed48b9b905 100644
--- a/crates/nu_plugin_polars/src/dataframe/series/string/strftime.rs
+++ b/crates/nu_plugin_polars/src/dataframe/series/string/strftime.rs
@@ -1,4 +1,9 @@
-use crate::{values::CustomValueSupport, PolarsPlugin};
+use crate::{
+ values::{
+ cant_convert_err, CustomValueSupport, NuExpression, PolarsPluginObject, PolarsPluginType,
+ },
+ PolarsPlugin,
+};
use super::super::super::values::{Column, NuDataFrame};
@@ -26,34 +31,60 @@ impl PluginCommand for StrFTime {
fn signature(&self) -> Signature {
Signature::build(self.name())
.required("fmt", SyntaxShape::String, "Format rule")
- .input_output_type(
- Type::Custom("dataframe".into()),
- Type::Custom("dataframe".into()),
- )
+ .input_output_types(vec![
+ (
+ Type::Custom("expression".into()),
+ Type::Custom("expression".into()),
+ ),
+ (
+ Type::Custom("dataframe".into()),
+ Type::Custom("dataframe".into()),
+ ),
+ ])
.category(Category::Custom("dataframe".into()))
}
fn examples(&self) -> Vec {
- vec![Example {
- description: "Formats date",
- example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime --timezone 'UTC');
+ vec![
+ Example {
+ description: "Formats date column as a string",
+ example: r#"let date = '2020-08-04T16:39:18+00:00' | into datetime --timezone 'UTC';
+ let df = ([[a]; [$date]] | polars into-df);
+ let df2 = $df | polars with-column [(polars col a | polars strftime "%Y/%m/%d" | polars as b)] | polars collect;
+ $df2.b"#,
+ result: Some(
+ NuDataFrame::try_from_columns(
+ vec![Column::new(
+ "b".to_string(),
+ vec![Value::test_string("2020/08/04")],
+ )],
+ None,
+ )
+ .expect("simple df for test should not fail")
+ .into_value(Span::test_data()),
+ ),
+ },
+ Example {
+ description: "Formats date",
+ example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime --timezone 'UTC');
let df = ([$dt $dt] | polars into-df);
$df | polars strftime "%Y/%m/%d""#,
- result: Some(
- NuDataFrame::try_from_columns(
- vec![Column::new(
- "0".to_string(),
- vec![
- Value::test_string("2020/08/04"),
- Value::test_string("2020/08/04"),
- ],
- )],
- None,
- )
- .expect("simple df for test should not fail")
- .into_value(Span::test_data()),
- ),
- }]
+ result: Some(
+ NuDataFrame::try_from_columns(
+ vec![Column::new(
+ "0".to_string(),
+ vec![
+ Value::test_string("2020/08/04"),
+ Value::test_string("2020/08/04"),
+ ],
+ )],
+ None,
+ )
+ .expect("simple df for test should not fail")
+ .into_value(Span::test_data()),
+ ),
+ },
+ ]
}
fn run(
@@ -63,19 +94,45 @@ impl PluginCommand for StrFTime {
call: &EvaluatedCall,
input: PipelineData,
) -> Result {
- command(plugin, engine, call, input).map_err(LabeledError::from)
+ let value = input.into_value(call.head)?;
+ match PolarsPluginObject::try_from_value(plugin, &value)? {
+ PolarsPluginObject::NuDataFrame(df) => command_df(plugin, engine, call, df),
+ PolarsPluginObject::NuLazyFrame(lazy) => {
+ command_df(plugin, engine, call, lazy.collect(call.head)?)
+ }
+ PolarsPluginObject::NuExpression(expr) => command_expr(plugin, engine, call, expr),
+ _ => Err(cant_convert_err(
+ &value,
+ &[
+ PolarsPluginType::NuDataFrame,
+ PolarsPluginType::NuLazyFrame,
+ PolarsPluginType::NuExpression,
+ ],
+ )),
+ }
+ .map_err(LabeledError::from)
}
}
-fn command(
+fn command_expr(
plugin: &PolarsPlugin,
engine: &EngineInterface,
call: &EvaluatedCall,
- input: PipelineData,
+ expr: NuExpression,
+) -> Result {
+ let fmt: String = call.req(0)?;
+ let res: NuExpression = expr.into_polars().dt().strftime(&fmt).into();
+ res.to_pipeline_data(plugin, engine, call.head)
+}
+
+fn command_df(
+ plugin: &PolarsPlugin,
+ engine: &EngineInterface,
+ call: &EvaluatedCall,
+ df: NuDataFrame,
) -> Result {
let fmt: String = call.req(0)?;
- let df = NuDataFrame::try_from_pipeline_coerce(plugin, input, call.head)?;
let series = df.as_series(call.head)?;
let casted = series.datetime().map_err(|e| ShellError::GenericError {