extern/nushell
1
0
Fork 1
mirror of https://github.com/nushell/nushell.git synced 2025-06-20 18:08:36 +02:00
Code Issues Packages Projects Releases Wiki Activity

Introducing polars into-schema

Browse Source
This commit is contained in:
Jack Wright 2025-04-09 13:18:30 -07:00
parent b0f9cda9b5
commit b1c4fdfe1d
5 changed files with 105 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
crates/nu_plugin_polars/src/dataframe/command/core/mod.rs
View File

@ -13,6 +13,7 @@ mod to_dtype;
mod to_lazy;
mod to_nu;
mod to_repr;
mod to_schema;
pub use self::open::OpenDataFrame;
use crate::PolarsPlugin;
@ -42,5 +43,6 @@ pub(crate) fn core_commands() -> Vec<Box<dyn PluginCommand<Plugin = PolarsPlugin
Box::new(ToLazyFrame),
Box::new(ToRepr),
Box::new(to_dtype::ToDataType),
Box::new(to_schema::ToSchema),
]
}

4
crates/nu_plugin_polars/src/dataframe/command/core/open.rs
View File

@ -86,7 +86,7 @@ impl PluginCommand for OpenDataFrame {
)
.named(
"schema",
SyntaxShape::Record(vec![]),
SyntaxShape::Any,
r#"Polars Schema in format [{name: str}]. CSV, JSON, and JSONL files"#,
Some('s')
)
@ -103,7 +103,7 @@ impl PluginCommand for OpenDataFrame {
)
.named(
"hive-schema",
SyntaxShape::Record(vec![]),
SyntaxShape::Any,
r#"Hive schema in format [{name: str}]. Parquet and Arrow files"#,
None,
)

15
crates/nu_plugin_polars/src/dataframe/command/core/to_df.rs
View File

@ -34,7 +34,7 @@ impl PluginCommand for ToDataFrame {
Signature::build(self.name())
.named(
"schema",
SyntaxShape::Record(vec![]),
SyntaxShape::Any,
r#"Polars Schema in format [{name: str}]."#,
Some('s'),
)
@ -203,7 +203,18 @@ impl PluginCommand for ToDataFrame {
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),
}
},
Example {
description: "Use a predefined schama",
example: r#"let schema = {a: str, b: str}; [[a b]; [1 "foo"] [2 "bar"]] | polars into-df -s {a: str, b: str}"#,
result: Some(NuDataFrame::try_from_series_vec(vec![
Series::new("a".into(), ["1", "2"]),
Series::new("b".into(), ["foo", "bar"]),
], Span::test_data())
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),
},
]
}

39
crates/nu_plugin_polars/src/dataframe/command/core/to_lazy.rs
View File

@ -4,8 +4,11 @@ use crate::values::{NuDataFrame, NuLazyFrame};
use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand};
use nu_protocol::{
record, Category, Example, LabeledError, PipelineData, Signature, SyntaxShape, Type, Value,
record, Category, Example, LabeledError, PipelineData, Signature, Span, SyntaxShape, Type,
Value,
};
use polars::prelude::NamedFrom;
use polars::series::Series;
#[derive(Clone)]
pub struct ToLazyFrame;
@ -25,7 +28,7 @@ impl PluginCommand for ToLazyFrame {
Signature::build(self.name())
.named(
"schema",
SyntaxShape::Record(vec![]),
SyntaxShape::Any,
r#"Polars Schema in format [{name: str}]."#,
Some('s'),
)
@ -34,16 +37,28 @@ impl PluginCommand for ToLazyFrame {
}
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Takes a table and creates a lazyframe",
example: "[[a b];[1 2] [3 4]] | polars into-lazy",
result: None,
},
Example {
description: "Takes a table, creates a lazyframe, assigns column 'b' type str, displays the schema",
example: "[[a b];[1 2] [3 4]] | polars into-lazy --schema {b: str} | polars schema",
result: Some(Value::test_record(record! {"b" => Value::test_string("str")})),
},
vec![
Example {
description: "Takes a table and creates a lazyframe",
example: "[[a b];[1 2] [3 4]] | polars into-lazy",
result: None,
},
Example {
description: "Takes a table, creates a lazyframe, assigns column 'b' type str, displays the schema",
example: "[[a b];[1 2] [3 4]] | polars into-lazy --schema {b: str} | polars schema",
result: Some(Value::test_record(record! {"b" => Value::test_string("str")})),
},
Example {
description: "Use a predefined schama",
example: r#"let schema = {a: str, b: str}; [[a b]; [1 "foo"] [2 "bar"]] | polars into-lazy -s {a: str, b: str}"#,
result: Some(NuDataFrame::try_from_series_vec(vec![
Series::new("a".into(), ["1", "2"]),
Series::new("b".into(), ["foo", "bar"]),
], Span::test_data())
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),
},
]
}

61
crates/nu_plugin_polars/src/dataframe/command/core/to_schema.rs Normal file
View File

@ -0,0 +1,61 @@
use nu_plugin::PluginCommand;
use nu_protocol::{record, Category, Example, ShellError, Signature, Span, Type, Value};
use crate::{
values::{CustomValueSupport, NuSchema},
PolarsPlugin,
};
pub struct ToSchema;
impl PluginCommand for ToSchema {
type Plugin = PolarsPlugin;
fn name(&self) -> &str {
"polars into-schema"
}
fn description(&self) -> &str {
"Convert a value to a polars schema object"
}
fn signature(&self) -> Signature {
Signature::build(self.name())
.input_output_type(Type::Any, Type::Custom("schema".into()))
.category(Category::Custom("dataframe".into()))
}
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "Convert a record into a schema",
example: r#"{a: str, b: u8} | polars into-schema"#,
result: Some(Value::record(
record! {
"a" => Value::string("str", Span::test_data()),
"b" => Value::string("u8", Span::test_data()),
},
Span::test_data(),
)),
}]
}
fn run(
&self,
plugin: &Self::Plugin,
engine: &nu_plugin::EngineInterface,
call: &nu_plugin::EvaluatedCall,
input: nu_protocol::PipelineData,
) -> Result<nu_protocol::PipelineData, nu_protocol::LabeledError> {
command(plugin, engine, call, input).map_err(nu_protocol::LabeledError::from)
}
}
fn command(
plugin: &PolarsPlugin,
engine: &nu_plugin::EngineInterface,
call: &nu_plugin::EvaluatedCall,
input: nu_protocol::PipelineData,
) -> Result<nu_protocol::PipelineData, ShellError> {
NuSchema::try_from_pipeline(plugin, input, call.head)?
.to_pipeline_data(plugin, engine, call.head)
}