From 7d6f8b6640d92e69d8a11b8c72a1747d381ce711 Mon Sep 17 00:00:00 2001 From: Jack Wright Date: Tue, 8 Apr 2025 13:29:16 -0700 Subject: [PATCH] Work on NuDataType --- crates/nu_plugin_polars/src/cache/list.rs | 17 ++- .../src/dataframe/values/mod.rs | 10 +- .../src/dataframe/values/nu_dtype.rs | 18 --- .../dataframe/values/nu_dtype/custom_value.rs | 65 +++++++++++ .../src/dataframe/values/nu_dtype/mod.rs | 108 ++++++++++++++++++ .../src/dataframe/values/nu_schema.rs | 2 +- 6 files changed, 198 insertions(+), 22 deletions(-) delete mode 100644 crates/nu_plugin_polars/src/dataframe/values/nu_dtype.rs create mode 100644 crates/nu_plugin_polars/src/dataframe/values/nu_dtype/custom_value.rs create mode 100644 crates/nu_plugin_polars/src/dataframe/values/nu_dtype/mod.rs diff --git a/crates/nu_plugin_polars/src/cache/list.rs b/crates/nu_plugin_polars/src/cache/list.rs index c68f910a81..4afaf9fa23 100644 --- a/crates/nu_plugin_polars/src/cache/list.rs +++ b/crates/nu_plugin_polars/src/cache/list.rs @@ -132,7 +132,22 @@ impl PluginCommand for ListDF { "reference_count" => Value::int(value.reference_count as i64, call.head), }, call.head, - ))) + ))), + PolarsPluginObject::NuDataType(_) => Ok(Some(Value::record( + record! { + "key" => Value::string(key.to_string(), call.head), + "created" => Value::date(value.created, call.head), + "columns" => Value::nothing(call.head), + "rows" => Value::nothing(call.head), + "type" => Value::string("DataType", call.head), + "estimated_size" => Value::nothing(call.head), + "span_contents" => Value::string(span_contents, value.span), + "span_start" => Value::int(value.span.start as i64, call.head), + "span_end" => Value::int(value.span.end as i64, call.head), + "reference_count" => Value::int(value.reference_count as i64, call.head), + }, + call.head, + ))), } })?; let vals = vals.into_iter().flatten().collect(); diff --git a/crates/nu_plugin_polars/src/dataframe/values/mod.rs b/crates/nu_plugin_polars/src/dataframe/values/mod.rs index c3db22db19..02a1ab10a3 100644 --- a/crates/nu_plugin_polars/src/dataframe/values/mod.rs +++ b/crates/nu_plugin_polars/src/dataframe/values/mod.rs @@ -1,12 +1,12 @@ mod file_type; mod nu_dataframe; +mod nu_dtype; mod nu_expression; mod nu_lazyframe; mod nu_lazygroupby; mod nu_schema; mod nu_when; pub mod utils; -mod nu_dtype; use crate::{Cacheable, PolarsPlugin}; use nu_plugin::EngineInterface; @@ -18,11 +18,11 @@ use uuid::Uuid; pub use file_type::PolarsFileType; pub use nu_dataframe::{Axis, Column, NuDataFrame, NuDataFrameCustomValue}; +pub use nu_dtype::NuDataType; pub use nu_expression::{NuExpression, NuExpressionCustomValue}; pub use nu_lazyframe::{NuLazyFrame, NuLazyFrameCustomValue}; pub use nu_lazygroupby::{NuLazyGroupBy, NuLazyGroupByCustomValue}; pub use nu_schema::{str_to_dtype, NuSchema}; -pub use nu_dtype::NuDataType; pub use nu_when::{NuWhen, NuWhenCustomValue, NuWhenType}; #[derive(Debug, Clone)] @@ -33,6 +33,7 @@ pub enum PolarsPluginType { NuLazyGroupBy, NuWhen, NuPolarsTestData, + NuDataType, } impl fmt::Display for PolarsPluginType { @@ -44,6 +45,7 @@ impl fmt::Display for PolarsPluginType { Self::NuLazyGroupBy => write!(f, "NuLazyGroupBy"), Self::NuWhen => write!(f, "NuWhen"), Self::NuPolarsTestData => write!(f, "NuPolarsTestData"), + Self::NuDataType => write!(f, "NuDataType"), } } } @@ -56,6 +58,7 @@ pub enum PolarsPluginObject { NuLazyGroupBy(NuLazyGroupBy), NuWhen(NuWhen), NuPolarsTestData(Uuid, String), + NuDataType(NuDataType), } impl PolarsPluginObject { @@ -104,6 +107,7 @@ impl PolarsPluginObject { Self::NuLazyGroupBy(_) => PolarsPluginType::NuLazyGroupBy, Self::NuWhen(_) => PolarsPluginType::NuWhen, Self::NuPolarsTestData(_, _) => PolarsPluginType::NuPolarsTestData, + Self::NuDataType(_) => PolarsPluginType::NuDataType, } } @@ -115,6 +119,7 @@ impl PolarsPluginObject { PolarsPluginObject::NuLazyGroupBy(lg) => lg.id, PolarsPluginObject::NuWhen(w) => w.id, PolarsPluginObject::NuPolarsTestData(id, _) => *id, + PolarsPluginObject::NuDataType(dt) => dt.id, } } @@ -128,6 +133,7 @@ impl PolarsPluginObject { PolarsPluginObject::NuPolarsTestData(id, s) => { Value::string(format!("{id}:{s}"), Span::test_data()) } + PolarsPluginObject::NuDataType(dt) => dt.into_value(span), } } diff --git a/crates/nu_plugin_polars/src/dataframe/values/nu_dtype.rs b/crates/nu_plugin_polars/src/dataframe/values/nu_dtype.rs deleted file mode 100644 index 7ece475319..0000000000 --- a/crates/nu_plugin_polars/src/dataframe/values/nu_dtype.rs +++ /dev/null @@ -1,18 +0,0 @@ -use nu_protocol::{ShellError, Span}; -use polars::prelude::{DataType, UnknownKind}; - -#[derive(Debug, Clone)] -pub struct NuDataType { - dtype: DataType, -} - -impl NuDataType { - pub fn new(dtype: DataType) -> Self { - Self { dtype } - } - - pub fn to_polars(&self) -> DataType { - self.dtype.clone() - } -} - diff --git a/crates/nu_plugin_polars/src/dataframe/values/nu_dtype/custom_value.rs b/crates/nu_plugin_polars/src/dataframe/values/nu_dtype/custom_value.rs new file mode 100644 index 0000000000..ae5e4113b3 --- /dev/null +++ b/crates/nu_plugin_polars/src/dataframe/values/nu_dtype/custom_value.rs @@ -0,0 +1,65 @@ +use nu_protocol::{CustomValue, ShellError, Span, Value}; +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +use crate::values::{CustomValueSupport, PolarsPluginCustomValue}; + +use super::NuDataType; + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct NuDataTypeCustomValue { + pub id: Uuid, + #[serde(skip)] + pub datatype: Option, +} + +#[typetag::serde] +impl CustomValue for NuDataTypeCustomValue { + fn clone_value(&self, span: nu_protocol::Span) -> Value { + Value::custom(Box::new(self.clone()), span) + } + + fn type_name(&self) -> String { + "NuDataFrame".into() + } + + fn to_base_value(&self, span: Span) -> Result { + Ok(Value::string( + "NuDataFrameValue: custom_value_to_base_value should've been called", + span, + )) + } + + fn as_mut_any(&mut self) -> &mut dyn std::any::Any { + self + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn notify_plugin_on_drop(&self) -> bool { + true + } +} + +impl PolarsPluginCustomValue for NuDataTypeCustomValue { + type PolarsPluginObjectType = NuDataType; + + fn id(&self) -> &Uuid { + &self.id + } + + fn internal(&self) -> &Option { + &self.datatype + } + + fn custom_value_to_base_value( + &self, + plugin: &crate::PolarsPlugin, + _engine: &nu_plugin::EngineInterface, + ) -> Result { + let dtype = NuDataType::try_from_custom_value(plugin, self)?; + dtype.base_value(Span::unknown()) + } +} diff --git a/crates/nu_plugin_polars/src/dataframe/values/nu_dtype/mod.rs b/crates/nu_plugin_polars/src/dataframe/values/nu_dtype/mod.rs new file mode 100644 index 0000000000..8a470a51b3 --- /dev/null +++ b/crates/nu_plugin_polars/src/dataframe/values/nu_dtype/mod.rs @@ -0,0 +1,108 @@ +mod custom_value; + +use custom_value::NuDataTypeCustomValue; +use nu_protocol::{ShellError, Span, Value}; +use polars::prelude::DataType; +use uuid::Uuid; + +use crate::Cacheable; + +use super::{ + nu_schema::dtype_to_value, str_to_dtype, CustomValueSupport, PolarsPluginObject, + PolarsPluginType, +}; + +#[derive(Debug, Clone)] +pub struct NuDataType { + pub id: uuid::Uuid, + dtype: DataType, +} + +impl NuDataType { + pub fn new(dtype: DataType) -> Self { + Self { + id: uuid::Uuid::new_v4(), + dtype, + } + } + + pub fn new_with_str(dtype: &str, span: Span) -> Result { + let dtype = str_to_dtype(dtype, span)?; + Ok(Self { + id: uuid::Uuid::new_v4(), + dtype, + }) + } + + pub fn to_polars(&self) -> DataType { + self.dtype.clone() + } +} + +impl TryFrom<&Value> for NuDataType { + type Error = ShellError; + + fn try_from(value: &Value) -> Result { + match value { + Value::String { val, internal_span } => NuDataType::new_with_str(val, *internal_span), + _ => Err(ShellError::GenericError { + error: format!("Unsupported value: {:?}", value), + msg: "".into(), + span: Some(value.span()), + help: None, + inner: vec![], + }), + } + } +} + +impl From for Value { + fn from(nu_dtype: NuDataType) -> Self { + Value::String { + val: nu_dtype.dtype.to_string(), + internal_span: Span::unknown(), + } + } +} + +impl Cacheable for NuDataType { + fn cache_id(&self) -> &Uuid { + &self.id + } + + fn to_cache_value(&self) -> Result { + Ok(PolarsPluginObject::NuDataType(self.clone())) + } + + fn from_cache_value(cv: super::PolarsPluginObject) -> Result { + match cv { + PolarsPluginObject::NuDataType(dt) => Ok(dt), + _ => Err(ShellError::GenericError { + error: "Cache value is not a dataframe".into(), + msg: "".into(), + span: None, + help: None, + inner: vec![], + }), + } + } +} + +impl CustomValueSupport for NuDataType { + type CV = NuDataTypeCustomValue; + + fn get_type_static() -> super::PolarsPluginType { + PolarsPluginType::NuDataType + } + + fn custom_value(self) -> Self::CV { + NuDataTypeCustomValue { + id: self.id, + datatype: Some(self), + } + } + + fn base_value(self, span: Span) -> Result { + Ok(dtype_to_value(&self.dtype, span)) + } +} diff --git a/crates/nu_plugin_polars/src/dataframe/values/nu_schema.rs b/crates/nu_plugin_polars/src/dataframe/values/nu_schema.rs index 1e2ae4723a..51f8179261 100644 --- a/crates/nu_plugin_polars/src/dataframe/values/nu_schema.rs +++ b/crates/nu_plugin_polars/src/dataframe/values/nu_schema.rs @@ -55,7 +55,7 @@ fn fields_to_value(fields: impl Iterator, span: Span) -> Value { Value::record(record, Span::unknown()) } -fn dtype_to_value(dtype: &DataType, span: Span) -> Value { +pub fn dtype_to_value(dtype: &DataType, span: Span) -> Value { match dtype { DataType::Struct(fields) => fields_to_value(fields.iter().cloned(), span), _ => Value::string(dtype.to_string().replace('[', "<").replace(']', ">"), span),