nu_plugin_polars 0.112.2

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

mod custom_value;

use core::fmt;
use nu_protocol::shell_error::generic::GenericError;
use nu_protocol::{ShellError, Value, record};
use polars::prelude::LazyGroupBy;
use std::sync::Arc;
use uuid::Uuid;

use crate::Cacheable;

pub use self::custom_value::NuLazyGroupByCustomValue;

use super::{CustomValueSupport, NuSchema, PolarsPluginObject, PolarsPluginType};

// Lazyframe wrapper for Nushell operations
// Polars LazyFrame is behind and Option to allow easy implementation of
// the Deserialize trait
#[derive(Clone)]
pub struct NuLazyGroupBy {
    pub id: Uuid,
    pub group_by: Arc<LazyGroupBy>,
    pub schema: NuSchema,
    pub from_eager: bool,
}

impl fmt::Debug for NuLazyGroupBy {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "NuLazyGroupBy")
    }
}

impl NuLazyGroupBy {
    pub fn new(group_by: LazyGroupBy, from_eager: bool, schema: NuSchema) -> Self {
        NuLazyGroupBy {
            id: Uuid::new_v4(),
            group_by: Arc::new(group_by),
            from_eager,
            schema,
        }
    }

    pub fn to_polars(&self) -> LazyGroupBy {
        (*self.group_by).clone()
    }
}

impl Cacheable for NuLazyGroupBy {
    fn cache_id(&self) -> &Uuid {
        &self.id
    }

    fn to_cache_value(&self) -> Result<PolarsPluginObject, ShellError> {
        Ok(PolarsPluginObject::NuLazyGroupBy(self.clone()))
    }

    fn from_cache_value(cv: PolarsPluginObject) -> Result<Self, ShellError> {
        match cv {
            PolarsPluginObject::NuLazyGroupBy(df) => Ok(df),
            _ => Err(ShellError::Generic(GenericError::new_internal(
                "Cache value is not a group by",
                "",
            ))),
        }
    }
}

impl CustomValueSupport for NuLazyGroupBy {
    type CV = NuLazyGroupByCustomValue;

    fn custom_value(self) -> Self::CV {
        NuLazyGroupByCustomValue {
            id: self.id,
            groupby: Some(self),
        }
    }

    fn get_type_static() -> PolarsPluginType {
        PolarsPluginType::NuLazyGroupBy
    }

    fn base_value(self, span: nu_protocol::Span) -> Result<nu_protocol::Value, ShellError> {
        Ok(Value::record(
            record! {
                "LazyGroupBy" => Value::string("apply aggregation to complete execution plan", span)
            },
            span,
        ))
    }
}