nu_plugin_polars/dataframe/command/datetime/
get_second.rs

1use crate::{
2    PolarsPlugin,
3    values::{
4        CustomValueSupport, NuDataFrame, NuExpression, NuLazyFrame, PolarsPluginObject,
5        PolarsPluginType, cant_convert_err,
6    },
7};
8
9use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand};
10use nu_protocol::{Category, Example, LabeledError, PipelineData, ShellError, Signature, Span};
11use polars::{
12    prelude::{DatetimeMethods, IntoSeries, NamedFrom, col},
13    series::Series,
14};
15
16#[derive(Clone)]
17pub struct GetSecond;
18
19impl PluginCommand for GetSecond {
20    type Plugin = PolarsPlugin;
21
22    fn name(&self) -> &str {
23        "polars get-second"
24    }
25
26    fn description(&self) -> &str {
27        "Gets second from date."
28    }
29
30    fn signature(&self) -> Signature {
31        Signature::build(self.name())
32            .input_output_types(vec![
33                (
34                    PolarsPluginType::NuDataFrame.into(),
35                    PolarsPluginType::NuDataFrame.into(),
36                ),
37                (
38                    PolarsPluginType::NuLazyFrame.into(),
39                    PolarsPluginType::NuLazyFrame.into(),
40                ),
41                (
42                    PolarsPluginType::NuExpression.into(),
43                    PolarsPluginType::NuExpression.into(),
44                ),
45            ])
46            .category(Category::Custom("dataframe".into()))
47    }
48
49    fn examples(&self) -> Vec<Example<'_>> {
50        vec![
51            Example {
52                description: "Returns second from a date",
53                example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime --timezone 'UTC');
54    let df = ([$dt $dt] | polars into-df);
55    $df | polars get-second"#,
56                result: Some(
57                    NuDataFrame::try_from_series(
58                        Series::new("0".into(), &[18i8, 18]),
59                        Span::test_data(),
60                    )
61                    .expect("simple df for test should not fail")
62                    .into_value(Span::test_data()),
63                ),
64            },
65            Example {
66                description: "Returns second from a date in an expression",
67                example: r#"let dt = ('2020-08-04T16:39:18+00:00' | into datetime --timezone 'UTC');
68    let df = ([$dt $dt] | polars into-df);
69    $df | polars select (polars col 0 | polars get-second)"#,
70                result: Some(
71                    NuDataFrame::try_from_series(
72                        Series::new("0".into(), &[18i8, 18]),
73                        Span::test_data(),
74                    )
75                    .expect("simple df for test should not fail")
76                    .into_value(Span::test_data()),
77                ),
78            },
79        ]
80    }
81
82    fn run(
83        &self,
84        plugin: &Self::Plugin,
85        engine: &EngineInterface,
86        call: &EvaluatedCall,
87        input: PipelineData,
88    ) -> Result<PipelineData, LabeledError> {
89        let metadata = input.metadata();
90        command(plugin, engine, call, input)
91            .map_err(LabeledError::from)
92            .map(|pd| pd.set_metadata(metadata))
93    }
94}
95
96fn command(
97    plugin: &PolarsPlugin,
98    engine: &EngineInterface,
99    call: &EvaluatedCall,
100    input: PipelineData,
101) -> Result<PipelineData, ShellError> {
102    let value = input.into_value(call.head)?;
103
104    match PolarsPluginObject::try_from_value(plugin, &value)? {
105        PolarsPluginObject::NuLazyFrame(lazy) => command_lazy(plugin, engine, call, lazy),
106        PolarsPluginObject::NuDataFrame(df) => command_eager(plugin, engine, call, df),
107        PolarsPluginObject::NuExpression(expr) => {
108            let res: NuExpression = expr.into_polars().dt().second().into();
109            res.to_pipeline_data(plugin, engine, call.head)
110        }
111        _ => Err(cant_convert_err(
112            &value,
113            &[
114                PolarsPluginType::NuDataFrame,
115                PolarsPluginType::NuLazyFrame,
116                PolarsPluginType::NuExpression,
117            ],
118        )),
119    }
120}
121
122fn command_lazy(
123    plugin: &PolarsPlugin,
124    engine: &EngineInterface,
125    call: &EvaluatedCall,
126    lazy: NuLazyFrame,
127) -> Result<PipelineData, ShellError> {
128    NuLazyFrame::new(false, lazy.to_polars().select([col("*").dt().second()]))
129        .to_pipeline_data(plugin, engine, call.head)
130}
131
132fn command_eager(
133    plugin: &PolarsPlugin,
134    engine: &EngineInterface,
135    call: &EvaluatedCall,
136    df: NuDataFrame,
137) -> Result<PipelineData, ShellError> {
138    let series = df.as_series(call.head)?;
139
140    let casted = series.datetime().map_err(|e| ShellError::GenericError {
141        error: "Error casting to datetime type".into(),
142        msg: e.to_string(),
143        span: Some(call.head),
144        help: None,
145        inner: vec![],
146    })?;
147
148    let res = casted.second().into_series();
149
150    let df = NuDataFrame::try_from_series_vec(vec![res], call.head)?;
151    df.to_pipeline_data(plugin, engine, call.head)
152}
153
154#[cfg(test)]
155mod test {
156    use super::*;
157    use crate::test::test_polars_plugin_command_with_decls;
158    use nu_command::IntoDatetime;
159
160    #[test]
161    fn test_examples() -> Result<(), ShellError> {
162        test_polars_plugin_command_with_decls(&GetSecond, vec![Box::new(IntoDatetime)])
163    }
164}
nu_plugin_polars/dataframe/command/datetime/get_second.rs

nu_plugin_polars/dataframe/command/datetime/
get_second.rs
