nu_plugin_polars 0.112.0

Nushell dataframe plugin commands based on polars.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

use crate::{
    PolarsPlugin,
    dataframe::values::{NuDataFrame, NuExpression},
    values::{CustomValueSupport, PolarsPluginObject, PolarsPluginType, cant_convert_err},
};
use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand};
use nu_protocol::{
    Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Value,
    shell_error::generic::GenericError,
};
use polars::df;

#[derive(Clone)]
pub struct Over;

impl PluginCommand for Over {
    type Plugin = PolarsPlugin;

    fn name(&self) -> &str {
        "polars over"
    }

    fn description(&self) -> &str {
        "Compute expressions over a window group defined by partition expressions."
    }

    fn signature(&self) -> Signature {
        Signature::build(self.name())
            .rest(
                "partition by expressions",
                SyntaxShape::Any,
                "Expression(s) that define the partition window.",
            )
            .input_output_type(
                PolarsPluginType::NuExpression.into(),
                PolarsPluginType::NuExpression.into(),
            )
            .category(Category::Custom("lazyframe".into()))
    }

    fn examples(&self) -> Vec<Example<'_>> {
        vec![
            Example {
                description: "Compute expression over an aggregation window",
                example: r#"[[a b]; [x 2] [x 4] [y 6] [y 4]]
        | polars into-lazy
        | polars select a (polars col b | polars cumulative sum | polars over a | polars as cum_b)
        | polars collect"#,
                result: Some(
                    NuDataFrame::from(
                        df!(
                            "a" => &["x", "x", "y", "y"],
                            "cum_b" => &[2, 6, 6, 10]
                        )
                        .expect("should not fail"),
                    )
                    .into_value(Span::test_data()),
                ),
            },
            Example {
                description: "Compute expression over an aggregation window where partitions are defined by expressions",
                example: r#"[[a b]; [x 2] [X 4] [Y 6] [y 4]]
        | polars into-lazy
        | polars select a (polars col b | polars cumulative sum | polars over (polars col a | polars lowercase) | polars as cum_b)
        | polars collect"#,
                result: Some(
                    NuDataFrame::from(
                        df!(
                            "a" => &["x", "X", "Y", "y"],
                            "cum_b" => &[2, 6, 6, 10]
                        )
                        .expect("should not fail"),
                    )
                    .into_value(Span::test_data()),
                ),
            },
        ]
    }

    fn run(
        &self,
        plugin: &Self::Plugin,
        engine: &EngineInterface,
        call: &EvaluatedCall,
        mut input: PipelineData,
    ) -> Result<PipelineData, LabeledError> {
        let vals: Vec<Value> = call.rest(0)?;
        let expr_value = Value::list(vals, call.head);
        let expressions = NuExpression::extract_exprs(plugin, expr_value)?;

        let metadata = input.take_metadata();
        let input_value = input.into_value(call.head)?;

        match PolarsPluginObject::try_from_value(plugin, &input_value)? {
            PolarsPluginObject::NuExpression(expr) => {
                let expr: NuExpression = expr
                    .into_polars()
                    .over_with_options(Some(expressions), None, Default::default())
                    .map_err(|e| {
                        ShellError::Generic(GenericError::new(
                            format!("Error applying over expression: {e}"),
                            "",
                            call.head,
                        ))
                    })?
                    .into();
                expr.to_pipeline_data(plugin, engine, call.head)
            }
            _ => Err(cant_convert_err(
                &input_value,
                &[PolarsPluginType::NuExpression],
            )),
        }
        .map_err(LabeledError::from)
        .map(|pd| pd.set_metadata(metadata))
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::test::test_polars_plugin_command;
    use nu_protocol::ShellError;

    #[test]
    fn test_examples() -> Result<(), ShellError> {
        test_polars_plugin_command(&Over)
    }
}