diff --git a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_day.rs b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_day.rs index ca9c365cd1..cc5466a200 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_day.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_day.rs @@ -1,13 +1,17 @@ -use crate::{values::CustomValueSupport, PolarsPlugin}; - -use super::super::super::values::NuDataFrame; +use crate::{ + values::{ + cant_convert_err, CustomValueSupport, NuDataFrame, NuExpression, NuLazyFrame, + PolarsPluginObject, PolarsPluginType, + }, + PolarsPlugin, +}; use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; use nu_protocol::{ Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, Type, }; use polars::{ - prelude::{DatetimeMethods, IntoSeries, NamedFrom}, + prelude::{col, DatetimeMethods, IntoSeries, NamedFrom}, series::Series, }; @@ -27,25 +31,50 @@ impl PluginCommand for GetDay { fn signature(&self) -> Signature { Signature::build(self.name()) - .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".into())) } fn examples(&self) -> Vec { - vec![Example { - description: "Returns day from a date", - example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime --timezone 'UTC'); + vec![ + Example { + description: "Returns day from a 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 get-day"#, - result: Some( - NuDataFrame::try_from_series(Series::new("0".into(), &[4i8, 4]), Span::test_data()) + result: Some( + NuDataFrame::try_from_series( + Series::new("0".into(), &[4i8, 4]), + Span::test_data(), + ) .expect("simple df for test should not fail") .into_value(Span::test_data()), - ), - }] + ), + }, + Example { + description: "Returns day from a date in an expression", + example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime --timezone 'UTC'); + let df = ([$dt $dt] | polars into-df); + $df | polars select (polars col 0 | polars get-day)"#, + result: Some( + NuDataFrame::try_from_series( + Series::new("0".into(), &[4i8, 4]), + Span::test_data(), + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + ] } fn run( @@ -73,7 +102,42 @@ fn command( call: &EvaluatedCall, input: PipelineData, ) -> Result { - let df = NuDataFrame::try_from_pipeline_coerce(plugin, input, call.head)?; + let value = input.into_value(call.head)?; + + match PolarsPluginObject::try_from_value(plugin, &value)? { + PolarsPluginObject::NuLazyFrame(lazy) => command_lazy(plugin, engine, call, lazy), + PolarsPluginObject::NuDataFrame(df) => command_eager(plugin, engine, call, df), + PolarsPluginObject::NuExpression(expr) => { + let res: NuExpression = expr.into_polars().dt().day().into(); + res.to_pipeline_data(plugin, engine, call.head) + } + _ => Err(cant_convert_err( + &value, + &[ + PolarsPluginType::NuDataFrame, + PolarsPluginType::NuLazyFrame, + PolarsPluginType::NuExpression, + ], + )), + } +} + +fn command_lazy( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + lazy: NuLazyFrame, +) -> Result { + NuLazyFrame::new(false, lazy.to_polars().select([col("*").dt().day()])) + .to_pipeline_data(plugin, engine, call.head) +} + +fn command_eager( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + df: NuDataFrame, +) -> Result { let series = df.as_series(call.head)?; let casted = series.datetime().map_err(|e| ShellError::GenericError { diff --git a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_hour.rs b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_hour.rs index b924df973d..44db47acc5 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_hour.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_hour.rs @@ -1,4 +1,10 @@ -use crate::{values::CustomValueSupport, PolarsPlugin}; +use crate::{ + values::{ + cant_convert_err, CustomValueSupport, NuExpression, NuLazyFrame, PolarsPluginObject, + PolarsPluginType, + }, + PolarsPlugin, +}; use super::super::super::values::NuDataFrame; @@ -7,7 +13,7 @@ use nu_protocol::{ Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, Type, }; use polars::{ - prelude::{DatetimeMethods, IntoSeries, NamedFrom}, + prelude::{col, DatetimeMethods, IntoSeries, NamedFrom}, series::Series, }; @@ -22,33 +28,69 @@ impl PluginCommand for GetHour { } fn description(&self) -> &str { - "Gets hour from date." + "Gets hour from datetime." } fn signature(&self) -> Signature { Signature::build(self.name()) - .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".into())) } fn examples(&self) -> Vec { - vec![Example { - description: "Returns hour from a date", - example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime --timezone 'UTC'); + vec![ + Example { + description: "Returns hour from a 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 get-hour"#, - result: Some( - NuDataFrame::try_from_series( - Series::new("0".into(), &[16i8, 16]), - Span::test_data(), - ) - .expect("simple df for test should not fail") - .into_value(Span::test_data()), - ), - }] + result: Some( + NuDataFrame::try_from_series( + Series::new("0".into(), &[16i8, 16]), + Span::test_data(), + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + Example { + description: "Returns hour from a date in a lazyframe", + example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime --timezone 'UTC'); + let df = ([$dt $dt] | polars into-lazy); + $df | polars get-hour"#, + result: Some( + NuDataFrame::try_from_series( + Series::new("0".into(), &[16i8, 16]), + Span::test_data(), + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + Example { + description: "Returns hour from a date in an expression", + example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime --timezone 'UTC'); + let df = ([$dt $dt] | polars into-df); + $df | polars select (polars col 0 | polars get-hour)"#, + result: Some( + NuDataFrame::try_from_series( + Series::new("0".into(), &[16i8, 16]), + Span::test_data(), + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + ] } fn run( @@ -68,7 +110,42 @@ fn command( call: &EvaluatedCall, input: PipelineData, ) -> Result { - let df = NuDataFrame::try_from_pipeline_coerce(plugin, input, call.head)?; + let value = input.into_value(call.head)?; + + match PolarsPluginObject::try_from_value(plugin, &value)? { + PolarsPluginObject::NuLazyFrame(lazy) => command_lazy(plugin, engine, call, lazy), + PolarsPluginObject::NuDataFrame(df) => command_eager(plugin, engine, call, df), + PolarsPluginObject::NuExpression(expr) => { + let res: NuExpression = expr.into_polars().dt().hour().into(); + res.to_pipeline_data(plugin, engine, call.head) + } + _ => Err(cant_convert_err( + &value, + &[ + PolarsPluginType::NuDataFrame, + PolarsPluginType::NuLazyFrame, + PolarsPluginType::NuExpression, + ], + )), + } +} + +fn command_lazy( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + lazy: NuLazyFrame, +) -> Result { + NuLazyFrame::new(false, lazy.to_polars().select([col("*").dt().hour()])) + .to_pipeline_data(plugin, engine, call.head) +} + +fn command_eager( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + df: NuDataFrame, +) -> Result { let series = df.as_series(call.head)?; let casted = series.datetime().map_err(|e| ShellError::GenericError { diff --git a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_minute.rs b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_minute.rs index aac0507dff..f52b0b5033 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_minute.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_minute.rs @@ -1,13 +1,19 @@ -use crate::{values::CustomValueSupport, PolarsPlugin}; -use polars::{prelude::NamedFrom, series::Series}; - -use super::super::super::values::NuDataFrame; +use crate::{ + values::{ + cant_convert_err, CustomValueSupport, NuDataFrame, NuExpression, NuLazyFrame, + PolarsPluginObject, PolarsPluginType, + }, + PolarsPlugin, +}; use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; use nu_protocol::{ Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, Type, }; -use polars::prelude::{DatetimeMethods, IntoSeries}; +use polars::{ + prelude::{col, DatetimeMethods, IntoSeries, NamedFrom}, + series::Series, +}; #[derive(Clone)] pub struct GetMinute; @@ -25,28 +31,50 @@ impl PluginCommand for GetMinute { fn signature(&self) -> Signature { Signature::build(self.name()) - .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".into())) } fn examples(&self) -> Vec { - vec![Example { - description: "Returns minute from a date", - example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime --timezone 'UTC'); + vec![ + Example { + description: "Returns minute from a 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 get-minute"#, - result: Some( - NuDataFrame::try_from_series( - Series::new("0".into(), &[39i8, 39]), - Span::test_data(), - ) - .expect("simple df for test should not fail") - .into_value(Span::test_data()), - ), - }] + result: Some( + NuDataFrame::try_from_series( + Series::new("0".into(), &[39i8, 39]), + Span::test_data(), + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + Example { + description: "Returns minute from a 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 select (polars col 0 | polars get-minute)"#, + result: Some( + NuDataFrame::try_from_series( + Series::new("0".into(), &[39i8, 39]), + Span::test_data(), + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + ] } fn run( @@ -66,7 +94,42 @@ fn command( call: &EvaluatedCall, input: PipelineData, ) -> Result { - let df = NuDataFrame::try_from_pipeline_coerce(plugin, input, call.head)?; + let value = input.into_value(call.head)?; + + match PolarsPluginObject::try_from_value(plugin, &value)? { + PolarsPluginObject::NuLazyFrame(lazy) => command_lazy(plugin, engine, call, lazy), + PolarsPluginObject::NuDataFrame(df) => command_eager(plugin, engine, call, df), + PolarsPluginObject::NuExpression(expr) => { + let res: NuExpression = expr.into_polars().dt().minute().into(); + res.to_pipeline_data(plugin, engine, call.head) + } + _ => Err(cant_convert_err( + &value, + &[ + PolarsPluginType::NuDataFrame, + PolarsPluginType::NuLazyFrame, + PolarsPluginType::NuExpression, + ], + )), + } +} + +fn command_lazy( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + lazy: NuLazyFrame, +) -> Result { + NuLazyFrame::new(false, lazy.to_polars().select([col("*").dt().minute()])) + .to_pipeline_data(plugin, engine, call.head) +} + +fn command_eager( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + df: NuDataFrame, +) -> Result { let series = df.as_series(call.head)?; let casted = series.datetime().map_err(|e| ShellError::GenericError { diff --git a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_month.rs b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_month.rs index 60d9400e51..d52d281148 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_month.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_month.rs @@ -1,13 +1,17 @@ -use crate::{values::CustomValueSupport, PolarsPlugin}; - -use super::super::super::values::NuDataFrame; +use crate::{ + values::{ + cant_convert_err, CustomValueSupport, NuDataFrame, NuExpression, NuLazyFrame, + PolarsPluginObject, PolarsPluginType, + }, + PolarsPlugin, +}; use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; use nu_protocol::{ Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, Type, }; use polars::{ - prelude::{DatetimeMethods, IntoSeries, NamedFrom}, + prelude::{col, DatetimeMethods, IntoSeries, NamedFrom}, series::Series, }; @@ -27,25 +31,50 @@ impl PluginCommand for GetMonth { fn signature(&self) -> Signature { Signature::build(self.name()) - .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".into())) } fn examples(&self) -> Vec { - vec![Example { - description: "Returns month from a date", - example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime --timezone 'UTC'); + vec![ + Example { + description: "Returns month from a 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 get-month"#, - result: Some( - NuDataFrame::try_from_series(Series::new("0".into(), &[8i8, 8]), Span::test_data()) + result: Some( + NuDataFrame::try_from_series( + Series::new("0".into(), &[8i8, 8]), + Span::test_data(), + ) .expect("simple df for test should not fail") .into_value(Span::test_data()), - ), - }] + ), + }, + Example { + description: "Returns month from a date in an expression", + example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime --timezone 'UTC'); + let df = ([$dt $dt] | polars into-df); + $df | polars select (polars col 0 | polars get-month)"#, + result: Some( + NuDataFrame::try_from_series( + Series::new("0".into(), &[8i8, 8]), + Span::test_data(), + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + ] } fn run( @@ -65,7 +94,42 @@ fn command( call: &EvaluatedCall, input: PipelineData, ) -> Result { - let df = NuDataFrame::try_from_pipeline_coerce(plugin, input, call.head)?; + let value = input.into_value(call.head)?; + + match PolarsPluginObject::try_from_value(plugin, &value)? { + PolarsPluginObject::NuLazyFrame(lazy) => command_lazy(plugin, engine, call, lazy), + PolarsPluginObject::NuDataFrame(df) => command_eager(plugin, engine, call, df), + PolarsPluginObject::NuExpression(expr) => { + let res: NuExpression = expr.into_polars().dt().month().into(); + res.to_pipeline_data(plugin, engine, call.head) + } + _ => Err(cant_convert_err( + &value, + &[ + PolarsPluginType::NuDataFrame, + PolarsPluginType::NuLazyFrame, + PolarsPluginType::NuExpression, + ], + )), + } +} + +fn command_lazy( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + lazy: NuLazyFrame, +) -> Result { + NuLazyFrame::new(false, lazy.to_polars().select([col("*").dt().month()])) + .to_pipeline_data(plugin, engine, call.head) +} + +fn command_eager( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + df: NuDataFrame, +) -> Result { let series = df.as_series(call.head)?; let casted = series.datetime().map_err(|e| ShellError::GenericError { diff --git a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_nanosecond.rs b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_nanosecond.rs index 1a683e92bd..8db9cf30ee 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_nanosecond.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_nanosecond.rs @@ -1,13 +1,17 @@ -use crate::{values::CustomValueSupport, PolarsPlugin}; - -use super::super::super::values::NuDataFrame; +use crate::{ + values::{ + cant_convert_err, CustomValueSupport, NuDataFrame, NuExpression, NuLazyFrame, + PolarsPluginObject, PolarsPluginType, + }, + PolarsPlugin, +}; use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; use nu_protocol::{ Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, Type, }; use polars::{ - prelude::{DatetimeMethods, IntoSeries, NamedFrom}, + prelude::{col, DatetimeMethods, IntoSeries, NamedFrom}, series::Series, }; @@ -27,28 +31,50 @@ impl PluginCommand for GetNanosecond { fn signature(&self) -> Signature { Signature::build(self.name()) - .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".into())) } fn examples(&self) -> Vec { - vec![Example { - description: "Returns nanosecond from a date", - example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime --timezone 'UTC'); + vec![ + Example { + description: "Returns nanosecond from a 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 get-nanosecond"#, - result: Some( - NuDataFrame::try_from_series( - Series::new("0".into(), &[0i32, 0]), - Span::test_data(), - ) - .expect("simple df for test should not fail") - .into_value(Span::test_data()), - ), - }] + result: Some( + NuDataFrame::try_from_series( + Series::new("0".into(), &[0i32, 0]), + Span::test_data(), + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + Example { + description: "Returns nanosecond from a 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 select (polars col 0 | polars get-nanosecond)"#, + result: Some( + NuDataFrame::try_from_series( + Series::new("0".into(), &[0i32, 0]), + Span::test_data(), + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + ] } fn run( @@ -68,7 +94,42 @@ fn command( call: &EvaluatedCall, input: PipelineData, ) -> Result { - let df = NuDataFrame::try_from_pipeline_coerce(plugin, input, call.head)?; + let value = input.into_value(call.head)?; + + match PolarsPluginObject::try_from_value(plugin, &value)? { + PolarsPluginObject::NuLazyFrame(lazy) => command_lazy(plugin, engine, call, lazy), + PolarsPluginObject::NuDataFrame(df) => command_eager(plugin, engine, call, df), + PolarsPluginObject::NuExpression(expr) => { + let res: NuExpression = expr.into_polars().dt().nanosecond().into(); + res.to_pipeline_data(plugin, engine, call.head) + } + _ => Err(cant_convert_err( + &value, + &[ + PolarsPluginType::NuDataFrame, + PolarsPluginType::NuLazyFrame, + PolarsPluginType::NuExpression, + ], + )), + } +} + +fn command_lazy( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + lazy: NuLazyFrame, +) -> Result { + NuLazyFrame::new(false, lazy.to_polars().select([col("*").dt().nanosecond()])) + .to_pipeline_data(plugin, engine, call.head) +} + +fn command_eager( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + df: NuDataFrame, +) -> Result { let series = df.as_series(call.head)?; let casted = series.datetime().map_err(|e| ShellError::GenericError { diff --git a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_ordinal.rs b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_ordinal.rs index 1c62c2b2ca..c49b3c74fc 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_ordinal.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_ordinal.rs @@ -1,13 +1,17 @@ -use crate::{values::CustomValueSupport, PolarsPlugin}; - -use super::super::super::values::NuDataFrame; +use crate::{ + values::{ + cant_convert_err, CustomValueSupport, NuDataFrame, NuExpression, NuLazyFrame, + PolarsPluginObject, PolarsPluginType, + }, + PolarsPlugin, +}; use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; use nu_protocol::{ Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, Type, }; use polars::{ - prelude::{DatetimeMethods, IntoSeries, NamedFrom}, + prelude::{col, DatetimeMethods, IntoSeries, NamedFrom}, series::Series, }; @@ -27,28 +31,50 @@ impl PluginCommand for GetOrdinal { fn signature(&self) -> Signature { Signature::build(self.name()) - .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".into())) } fn examples(&self) -> Vec { - vec![Example { - description: "Returns ordinal from a date", - example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime --timezone 'UTC'); + vec![ + Example { + description: "Returns ordinal from a 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 get-ordinal"#, - result: Some( - NuDataFrame::try_from_series( - Series::new("0".into(), &[217i16, 217]), - Span::test_data(), - ) - .expect("simple df for test should not fail") - .into_value(Span::test_data()), - ), - }] + result: Some( + NuDataFrame::try_from_series( + Series::new("0".into(), &[217i16, 217]), + Span::test_data(), + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + Example { + description: "Returns ordinal from a date in an expression", + example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime --timezone 'UTC'); + let df = ([$dt $dt] | polars into-df); + $df | polars select (polars col 0 | polars get-ordinal)"#, + result: Some( + NuDataFrame::try_from_series( + Series::new("0".into(), &[217i16, 217]), + Span::test_data(), + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + ] } fn run( @@ -68,7 +94,45 @@ fn command( call: &EvaluatedCall, input: PipelineData, ) -> Result { - let df = NuDataFrame::try_from_pipeline_coerce(plugin, input, call.head)?; + let value = input.into_value(call.head)?; + + match PolarsPluginObject::try_from_value(plugin, &value)? { + PolarsPluginObject::NuLazyFrame(lazy) => command_lazy(plugin, engine, call, lazy), + PolarsPluginObject::NuDataFrame(df) => command_eager(plugin, engine, call, df), + PolarsPluginObject::NuExpression(expr) => { + let res: NuExpression = expr.into_polars().dt().ordinal_day().into(); + res.to_pipeline_data(plugin, engine, call.head) + } + _ => Err(cant_convert_err( + &value, + &[ + PolarsPluginType::NuDataFrame, + PolarsPluginType::NuLazyFrame, + PolarsPluginType::NuExpression, + ], + )), + } +} + +fn command_lazy( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + lazy: NuLazyFrame, +) -> Result { + NuLazyFrame::new( + false, + lazy.to_polars().select([col("*").dt().ordinal_day()]), + ) + .to_pipeline_data(plugin, engine, call.head) +} + +fn command_eager( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + df: NuDataFrame, +) -> Result { let series = df.as_series(call.head)?; let casted = series.datetime().map_err(|e| ShellError::GenericError { diff --git a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_second.rs b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_second.rs index c80eefbe4b..4ffa7e7da1 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_second.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_second.rs @@ -1,13 +1,17 @@ -use crate::{values::CustomValueSupport, PolarsPlugin}; - -use super::super::super::values::NuDataFrame; +use crate::{ + values::{ + cant_convert_err, CustomValueSupport, NuDataFrame, NuExpression, NuLazyFrame, + PolarsPluginObject, PolarsPluginType, + }, + PolarsPlugin, +}; use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; use nu_protocol::{ Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, Type, }; use polars::{ - prelude::{DatetimeMethods, IntoSeries, NamedFrom}, + prelude::{col, DatetimeMethods, IntoSeries, NamedFrom}, series::Series, }; @@ -27,28 +31,50 @@ impl PluginCommand for GetSecond { fn signature(&self) -> Signature { Signature::build(self.name()) - .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".into())) } fn examples(&self) -> Vec { - vec![Example { - description: "Returns second from a date", - example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime --timezone 'UTC'); + vec![ + Example { + description: "Returns second from a 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 get-second"#, - result: Some( - NuDataFrame::try_from_series( - Series::new("0".into(), &[18i8, 18]), - Span::test_data(), - ) - .expect("simple df for test should not fail") - .into_value(Span::test_data()), - ), - }] + result: Some( + NuDataFrame::try_from_series( + Series::new("0".into(), &[18i8, 18]), + Span::test_data(), + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + Example { + description: "Returns second from a date in an expression", + example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime --timezone 'UTC'); + let df = ([$dt $dt] | polars into-df); + $df | polars select (polars col 0 | polars get-second)"#, + result: Some( + NuDataFrame::try_from_series( + Series::new("0".into(), &[18i8, 18]), + Span::test_data(), + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + ] } fn run( @@ -68,7 +94,42 @@ fn command( call: &EvaluatedCall, input: PipelineData, ) -> Result { - let df = NuDataFrame::try_from_pipeline_coerce(plugin, input, call.head)?; + let value = input.into_value(call.head)?; + + match PolarsPluginObject::try_from_value(plugin, &value)? { + PolarsPluginObject::NuLazyFrame(lazy) => command_lazy(plugin, engine, call, lazy), + PolarsPluginObject::NuDataFrame(df) => command_eager(plugin, engine, call, df), + PolarsPluginObject::NuExpression(expr) => { + let res: NuExpression = expr.into_polars().dt().second().into(); + res.to_pipeline_data(plugin, engine, call.head) + } + _ => Err(cant_convert_err( + &value, + &[ + PolarsPluginType::NuDataFrame, + PolarsPluginType::NuLazyFrame, + PolarsPluginType::NuExpression, + ], + )), + } +} + +fn command_lazy( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + lazy: NuLazyFrame, +) -> Result { + NuLazyFrame::new(false, lazy.to_polars().select([col("*").dt().second()])) + .to_pipeline_data(plugin, engine, call.head) +} + +fn command_eager( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + df: NuDataFrame, +) -> Result { let series = df.as_series(call.head)?; let casted = series.datetime().map_err(|e| ShellError::GenericError { diff --git a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_week.rs b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_week.rs index 7cde91af76..3d3802d9a0 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_week.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_week.rs @@ -1,13 +1,17 @@ -use crate::{values::CustomValueSupport, PolarsPlugin}; - -use super::super::super::values::NuDataFrame; +use crate::{ + values::{ + cant_convert_err, CustomValueSupport, NuDataFrame, NuExpression, NuLazyFrame, + PolarsPluginObject, PolarsPluginType, + }, + PolarsPlugin, +}; use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; use nu_protocol::{ Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, Type, }; use polars::{ - prelude::{DatetimeMethods, IntoSeries, NamedFrom}, + prelude::{col, DatetimeMethods, IntoSeries, NamedFrom}, series::Series, }; @@ -27,28 +31,50 @@ impl PluginCommand for GetWeek { fn signature(&self) -> Signature { Signature::build(self.name()) - .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".into())) } fn examples(&self) -> Vec { - vec![Example { - description: "Returns week from a date", - example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime --timezone 'UTC'); + vec![ + Example { + description: "Returns week from a 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 get-week"#, - result: Some( - NuDataFrame::try_from_series( - Series::new("0".into(), &[32i8, 32]), - Span::test_data(), - ) - .expect("simple df for test should not fail") - .into_value(Span::test_data()), - ), - }] + result: Some( + NuDataFrame::try_from_series( + Series::new("0".into(), &[32i8, 32]), + Span::test_data(), + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + Example { + description: "Returns week from a date in an expression", + example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime --timezone 'UTC'); + let df = ([$dt $dt] | polars into-df); + $df | polars select (polars col 0 | polars get-week)"#, + result: Some( + NuDataFrame::try_from_series( + Series::new("0".into(), &[32i8, 32]), + Span::test_data(), + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + ] } fn run( @@ -68,7 +94,42 @@ fn command( call: &EvaluatedCall, input: PipelineData, ) -> Result { - let df = NuDataFrame::try_from_pipeline_coerce(plugin, input, call.head)?; + let value = input.into_value(call.head)?; + + match PolarsPluginObject::try_from_value(plugin, &value)? { + PolarsPluginObject::NuLazyFrame(lazy) => command_lazy(plugin, engine, call, lazy), + PolarsPluginObject::NuDataFrame(df) => command_eager(plugin, engine, call, df), + PolarsPluginObject::NuExpression(expr) => { + let res: NuExpression = expr.into_polars().dt().week().into(); + res.to_pipeline_data(plugin, engine, call.head) + } + _ => Err(cant_convert_err( + &value, + &[ + PolarsPluginType::NuDataFrame, + PolarsPluginType::NuLazyFrame, + PolarsPluginType::NuExpression, + ], + )), + } +} + +fn command_lazy( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + lazy: NuLazyFrame, +) -> Result { + NuLazyFrame::new(false, lazy.to_polars().select([col("*").dt().week()])) + .to_pipeline_data(plugin, engine, call.head) +} + +fn command_eager( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + df: NuDataFrame, +) -> Result { let series = df.as_series(call.head)?; let casted = series.datetime().map_err(|e| ShellError::GenericError { diff --git a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_weekday.rs b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_weekday.rs index a480779c8a..73c2348a4e 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_weekday.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_weekday.rs @@ -1,13 +1,17 @@ -use crate::{values::CustomValueSupport, PolarsPlugin}; - -use super::super::super::values::NuDataFrame; +use crate::{ + values::{ + cant_convert_err, CustomValueSupport, NuDataFrame, NuExpression, NuLazyFrame, + PolarsPluginObject, PolarsPluginType, + }, + PolarsPlugin, +}; use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; use nu_protocol::{ Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, Type, }; use polars::{ - prelude::{DatetimeMethods, IntoSeries, NamedFrom}, + prelude::{col, DatetimeMethods, IntoSeries, NamedFrom}, series::Series, }; @@ -27,25 +31,50 @@ impl PluginCommand for GetWeekDay { fn signature(&self) -> Signature { Signature::build(self.name()) - .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".into())) } fn examples(&self) -> Vec { - vec![Example { - description: "Returns weekday from a date", - example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime --timezone 'UTC'); + vec![ + Example { + description: "Returns weekday from a 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 get-weekday"#, - result: Some( - NuDataFrame::try_from_series(Series::new("0".into(), &[2i8, 2]), Span::test_data()) + result: Some( + NuDataFrame::try_from_series( + Series::new("0".into(), &[2i8, 2]), + Span::test_data(), + ) .expect("simple df for test should not fail") .into_value(Span::test_data()), - ), - }] + ), + }, + Example { + description: "Returns weekday from a 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 select (polars col 0 | polars get-weekday)"#, + result: Some( + NuDataFrame::try_from_series( + Series::new("0".into(), &[2i8, 2]), + Span::test_data(), + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + ] } fn run( @@ -65,7 +94,42 @@ fn command( call: &EvaluatedCall, input: PipelineData, ) -> Result { - let df = NuDataFrame::try_from_pipeline_coerce(plugin, input, call.head)?; + let value = input.into_value(call.head)?; + + match PolarsPluginObject::try_from_value(plugin, &value)? { + PolarsPluginObject::NuLazyFrame(lazy) => command_lazy(plugin, engine, call, lazy), + PolarsPluginObject::NuDataFrame(df) => command_eager(plugin, engine, call, df), + PolarsPluginObject::NuExpression(expr) => { + let res: NuExpression = expr.into_polars().dt().weekday().into(); + res.to_pipeline_data(plugin, engine, call.head) + } + _ => Err(cant_convert_err( + &value, + &[ + PolarsPluginType::NuDataFrame, + PolarsPluginType::NuLazyFrame, + PolarsPluginType::NuExpression, + ], + )), + } +} + +fn command_lazy( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + lazy: NuLazyFrame, +) -> Result { + NuLazyFrame::new(false, lazy.to_polars().select([col("*").dt().weekday()])) + .to_pipeline_data(plugin, engine, call.head) +} + +fn command_eager( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + df: NuDataFrame, +) -> Result { let series = df.as_series(call.head)?; let casted = series.datetime().map_err(|e| ShellError::GenericError { diff --git a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_year.rs b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_year.rs index b057e95af2..16205bbbf9 100644 --- a/crates/nu_plugin_polars/src/dataframe/command/datetime/get_year.rs +++ b/crates/nu_plugin_polars/src/dataframe/command/datetime/get_year.rs @@ -1,13 +1,17 @@ -use crate::{values::CustomValueSupport, PolarsPlugin}; - -use super::super::super::values::NuDataFrame; +use crate::{ + values::{ + cant_convert_err, CustomValueSupport, NuDataFrame, NuExpression, NuLazyFrame, + PolarsPluginObject, PolarsPluginType, + }, + PolarsPlugin, +}; use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; use nu_protocol::{ Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, Type, }; use polars::{ - prelude::{DatetimeMethods, IntoSeries, NamedFrom}, + prelude::{col, DatetimeMethods, IntoSeries, NamedFrom}, series::Series, }; @@ -27,28 +31,50 @@ impl PluginCommand for GetYear { fn signature(&self) -> Signature { Signature::build(self.name()) - .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".into())) } fn examples(&self) -> Vec { - vec![Example { - description: "Returns year from a date", - example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime --timezone 'UTC'); + vec![ + Example { + description: "Returns year from a 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 get-year"#, - result: Some( - NuDataFrame::try_from_series( - Series::new("0".into(), &[2020i32, 2020]), - Span::test_data(), - ) - .expect("simple df for test should not fail") - .into_value(Span::test_data()), - ), - }] + result: Some( + NuDataFrame::try_from_series( + Series::new("0".into(), &[2020i32, 2020]), + Span::test_data(), + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + Example { + description: "Returns year from a date in an expression", + example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime --timezone 'UTC'); + let df = ([$dt $dt] | polars into-df); + $df | polars select (polars col 0 | polars get-year)"#, + result: Some( + NuDataFrame::try_from_series( + Series::new("0".into(), &[2020i32, 2020]), + Span::test_data(), + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, + ] } fn run( @@ -68,7 +94,42 @@ fn command( call: &EvaluatedCall, input: PipelineData, ) -> Result { - let df = NuDataFrame::try_from_pipeline_coerce(plugin, input, call.head)?; + let value = input.into_value(call.head)?; + + match PolarsPluginObject::try_from_value(plugin, &value)? { + PolarsPluginObject::NuLazyFrame(lazy) => command_lazy(plugin, engine, call, lazy), + PolarsPluginObject::NuDataFrame(df) => command_eager(plugin, engine, call, df), + PolarsPluginObject::NuExpression(expr) => { + let res: NuExpression = expr.into_polars().dt().year().into(); + res.to_pipeline_data(plugin, engine, call.head) + } + _ => Err(cant_convert_err( + &value, + &[ + PolarsPluginType::NuDataFrame, + PolarsPluginType::NuLazyFrame, + PolarsPluginType::NuExpression, + ], + )), + } +} + +fn command_lazy( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + lazy: NuLazyFrame, +) -> Result { + NuLazyFrame::new(false, lazy.to_polars().select([col("*").dt().year()])) + .to_pipeline_data(plugin, engine, call.head) +} + +fn command_eager( + plugin: &PolarsPlugin, + engine: &EngineInterface, + call: &EvaluatedCall, + df: NuDataFrame, +) -> Result { let series = df.as_series(call.head)?; let casted = series.datetime().map_err(|e| ShellError::GenericError {