Public/nushell
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

cleanup.

Browse Source
This commit is contained in:
Ben Yang 2024-07-25 11:36:55 +08:00
parent f0ef1fbe37
commit 554ee5102b
1 changed files with 27 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

42
crates/nu_plugin_polars/src/dataframe/eager/to_df.rs
View File

@ -8,9 +8,9 @@ use super::super::values::NuDataFrame;
use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand}; use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand};
use nu_protocol::{ use nu_protocol::{
Category, Example, LabeledError, PipelineData, Record, ShellError, Signature, Span, SyntaxShape, Type, Value Category, Example, LabeledError, PipelineData, Signature, Span,
SyntaxShape, Type, Value,
}; };
use nu_utils::SharedCow;
use polars::{ use polars::{
prelude::{AnyValue, DataType, Field, NamedFrom}, prelude::{AnyValue, DataType, Field, NamedFrom},
series::Series, series::Series,
@ -38,7 +38,11 @@ impl PluginCommand for ToDataFrame {
r#"Polars Schema in format [{name: str}]. CSV, JSON, and JSONL files"#, r#"Polars Schema in format [{name: str}]. CSV, JSON, and JSONL files"#,
Some('s'), Some('s'),
) )
.switch("as-columns", r#"When input shape is record of lists, treat each list as column values."#, Some('c')) .switch(
"as-columns",
r#"When input shape is record of lists, treat each list as column values."#,
Some('c'),
)
.input_output_type(Type::Any, Type::Custom("dataframe".into())) .input_output_type(Type::Any, Type::Custom("dataframe".into()))
.category(Category::Custom("dataframe".into())) .category(Category::Custom("dataframe".into()))
} }
@ -206,27 +210,35 @@ impl PluginCommand for ToDataFrame {
.transpose()?; .transpose()?;
let maybe_as_columns = call.has_flag("as-columns")?; let maybe_as_columns = call.has_flag("as-columns")?;
let df = if !maybe_as_columns { let df = if !maybe_as_columns {
NuDataFrame::try_from_iter(plugin, input.into_iter(), maybe_schema.clone())? NuDataFrame::try_from_iter(plugin, input.into_iter(), maybe_schema.clone())?
} else { } else {
match &input { match &input {
PipelineData::Value(Value::Record {val, ..}, _) => { PipelineData::Value(Value::Record { val, .. }, _) => {
let items: Result<Vec<(String, Vec<Value>)>,&str> = val.iter().map( let items: Result<Vec<(String, Vec<Value>)>, &str> = val
|(k, v)| match v.to_owned().into_list() { .iter()
.map(|(k, v)| match v.to_owned().into_list() {
Ok(v) => Ok((k.to_owned(), v)), Ok(v) => Ok((k.to_owned(), v)),
_ => Err("error") _ => Err("error"),
} })
).collect(); .collect();
match items { match items {
Ok(items) => { Ok(items) => {
let columns = items.iter().map(|(k, v)| Column::new(k.to_owned(), v.to_owned())).collect::<Vec<Column>>(); let columns = items
.iter()
.map(|(k, v)| Column::new(k.to_owned(), v.to_owned()))
.collect::<Vec<Column>>();
NuDataFrame::try_from_columns(columns, maybe_schema)? NuDataFrame::try_from_columns(columns, maybe_schema)?
}, }
Err(_) => NuDataFrame::try_from_iter(plugin, input.into_iter(), maybe_schema.clone())? Err(_) => NuDataFrame::try_from_iter(
plugin,
input.into_iter(),
maybe_schema.clone(),
)?,
} }
}, }
_ => NuDataFrame::try_from_iter(plugin, input.into_iter(), maybe_schema.clone())? _ => NuDataFrame::try_from_iter(plugin, input.into_iter(), maybe_schema.clone())?,
} }
}; };