extern/nushell
1
0
Fork 1
mirror of https://github.com/nushell/nushell.git synced 2025-06-24 11:51:29 +02:00
Code Issues Packages Projects Releases Wiki Activity

Allow NuDataType values to be passed in when creating schemas

Browse Source
This commit is contained in:
Jack Wright 2025-04-08 16:52:11 -07:00
parent f44bd9863a
commit b2509a8084
1 changed files with 17 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
crates/nu_plugin_polars/src/dataframe/values/nu_schema/mod.rs
View File

@ -9,7 +9,7 @@ use uuid::Uuid;
use crate::{Cacheable, PolarsPlugin}; use crate::{Cacheable, PolarsPlugin};
use super::{str_to_dtype, CustomValueSupport, PolarsPluginObject, PolarsPluginType}; use super::{str_to_dtype, CustomValueSupport, NuDataType, PolarsPluginObject, PolarsPluginType};
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub struct NuSchema { pub struct NuSchema {
@ -92,7 +92,7 @@ impl CustomValueSupport for NuSchema {
}) })
} }
} else { } else {
let schema = value_to_schema(value, Span::unknown())?; let schema = value_to_schema(plugin, value, Span::unknown())?;
Ok(Self::new(Arc::new(schema))) Ok(Self::new(Arc::new(schema)))
} }
} }
@ -117,22 +117,30 @@ pub fn dtype_to_value(dtype: &DataType, span: Span) -> Value {
} }
} }
fn value_to_schema(value: &Value, span: Span) -> Result<Schema, ShellError> { fn value_to_schema(plugin: &PolarsPlugin, value: &Value, span: Span) -> Result<Schema, ShellError> {
let fields = value_to_fields(value, span)?; let fields = value_to_fields(plugin, value, span)?;
let schema = Schema::from_iter(fields); let schema = Schema::from_iter(fields);
Ok(schema) Ok(schema)
} }
fn value_to_fields(value: &Value, span: Span) -> Result<Vec<Field>, ShellError> { fn value_to_fields(
plugin: &PolarsPlugin,
value: &Value,
span: Span,
) -> Result<Vec<Field>, ShellError> {
let fields = value let fields = value
.as_record()? .as_record()?
.into_iter() .into_iter()
.map(|(col, val)| match val { .map(|(col, val)| match val {
Value::Record { .. } => { Value::Record { .. } => {
let fields = value_to_fields(val, span)?; let fields = value_to_fields(plugin, val, span)?;
let dtype = DataType::Struct(fields); let dtype = DataType::Struct(fields);
Ok(Field::new(col.into(), dtype)) Ok(Field::new(col.into(), dtype))
} }
Value::Custom { .. } => {
let dtype = NuDataType::try_from_value(plugin, val)?;
Ok(Field::new(col.into(), dtype.to_polars()))
}
_ => { _ => {
let dtype = str_to_dtype(&val.coerce_string()?, span)?; let dtype = str_to_dtype(&val.coerce_string()?, span)?;
Ok(Field::new(col.into(), dtype)) Ok(Field::new(col.into(), dtype))
@ -151,6 +159,8 @@ mod test {
#[test] #[test]
fn test_value_to_schema() { fn test_value_to_schema() {
let plugin = PolarsPlugin::new_test_mode().expect("Failed to create plugin");
let address = record! { let address = record! {
"street" => Value::test_string("str"), "street" => Value::test_string("str"),
"city" => Value::test_string("str"), "city" => Value::test_string("str"),
@ -162,7 +172,7 @@ mod test {
"address" => Value::test_record(address) "address" => Value::test_record(address)
}); });
let schema = value_to_schema(&value, Span::unknown()).unwrap(); let schema = value_to_schema(&plugin, &value, Span::unknown()).unwrap();
let expected = Schema::from_iter(vec![ let expected = Schema::from_iter(vec![
Field::new("name".into(), DataType::String), Field::new("name".into(), DataType::String),
Field::new("age".into(), DataType::Int32), Field::new("age".into(), DataType::Int32),