mirror of
https://github.com/nushell/nushell
synced 2024-12-26 13:03:07 +00:00
Implemented polars unnest
(#14104)
# Description Provides the ability to decomes struct columns into seperate columns for each field: <img width="655" alt="Screenshot 2024-10-16 at 09 57 22" src="https://github.com/user-attachments/assets/6706bd36-8d38-4365-b58d-ba82f2d5ba9a"> # User-Facing Changes - provides a new command `polars unnest` for decomposing struct fields into separate columns.
This commit is contained in:
parent
e911ff4d67
commit
dba2f6e0f8
2 changed files with 147 additions and 0 deletions
|
@ -32,6 +32,7 @@ mod sort_by_expr;
|
||||||
pub mod sql_context;
|
pub mod sql_context;
|
||||||
pub mod sql_expr;
|
pub mod sql_expr;
|
||||||
mod take;
|
mod take;
|
||||||
|
mod unnest;
|
||||||
mod unpivot;
|
mod unpivot;
|
||||||
mod with_column;
|
mod with_column;
|
||||||
use filter::LazyFilter;
|
use filter::LazyFilter;
|
||||||
|
@ -109,5 +110,6 @@ pub(crate) fn data_commands() -> Vec<Box<dyn PluginCommand<Plugin = PolarsPlugin
|
||||||
Box::new(LazyFilter),
|
Box::new(LazyFilter),
|
||||||
Box::new(Shift),
|
Box::new(Shift),
|
||||||
Box::new(Unique),
|
Box::new(Unique),
|
||||||
|
Box::new(unnest::UnnestDF),
|
||||||
]
|
]
|
||||||
}
|
}
|
||||||
|
|
145
crates/nu_plugin_polars/src/dataframe/command/data/unnest.rs
Normal file
145
crates/nu_plugin_polars/src/dataframe/command/data/unnest.rs
Normal file
|
@ -0,0 +1,145 @@
|
||||||
|
use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand};
|
||||||
|
use nu_protocol::{
|
||||||
|
Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape,
|
||||||
|
};
|
||||||
|
use polars::df;
|
||||||
|
|
||||||
|
use crate::{
|
||||||
|
values::{CustomValueSupport, NuLazyFrame, PolarsPluginObject},
|
||||||
|
PolarsPlugin,
|
||||||
|
};
|
||||||
|
|
||||||
|
use crate::values::NuDataFrame;
|
||||||
|
|
||||||
|
#[derive(Clone)]
|
||||||
|
pub struct UnnestDF;
|
||||||
|
|
||||||
|
impl PluginCommand for UnnestDF {
|
||||||
|
type Plugin = PolarsPlugin;
|
||||||
|
|
||||||
|
fn name(&self) -> &str {
|
||||||
|
"polars unnest"
|
||||||
|
}
|
||||||
|
|
||||||
|
fn description(&self) -> &str {
|
||||||
|
"Decompose struct columns into separate columns for each of their fields. The new columns will be inserted into the dataframe at the location of the struct column."
|
||||||
|
}
|
||||||
|
|
||||||
|
fn signature(&self) -> Signature {
|
||||||
|
Signature::build(self.name())
|
||||||
|
.rest("cols", SyntaxShape::String, "columns to unnest")
|
||||||
|
.category(Category::Custom("dataframe".into()))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn examples(&self) -> Vec<Example> {
|
||||||
|
vec![
|
||||||
|
Example {
|
||||||
|
description: "Unnest a dataframe",
|
||||||
|
example: r#"[[id person]; [1 {name: "Bob", age: 36}] [2 {name: "Betty", age: 63}]]
|
||||||
|
| polars into-df -s {id: i64, person: {name: str, age: u8}}
|
||||||
|
| polars unnest person
|
||||||
|
| polars get id name age
|
||||||
|
| polars sort-by id"#,
|
||||||
|
result: Some(
|
||||||
|
NuDataFrame::from(
|
||||||
|
df!(
|
||||||
|
"id" => [1, 2],
|
||||||
|
"name" => ["Bob", "Betty"],
|
||||||
|
"age" => [36, 63]
|
||||||
|
)
|
||||||
|
.expect("Should be able to create a simple dataframe"),
|
||||||
|
)
|
||||||
|
.into_value(Span::test_data()),
|
||||||
|
),
|
||||||
|
},
|
||||||
|
Example {
|
||||||
|
description: "Unnest a lazy dataframe",
|
||||||
|
example: r#"[[id person]; [1 {name: "Bob", age: 36}] [2 {name: "Betty", age: 63}]]
|
||||||
|
| polars into-df -s {id: i64, person: {name: str, age: u8}}
|
||||||
|
| polars into-lazy
|
||||||
|
| polars unnest person
|
||||||
|
| polars select (polars col id) (polars col name) (polars col age)
|
||||||
|
| polars collect
|
||||||
|
| polars sort-by id"#,
|
||||||
|
result: Some(
|
||||||
|
NuDataFrame::from(
|
||||||
|
df!(
|
||||||
|
"id" => [1, 2],
|
||||||
|
"name" => ["Bob", "Betty"],
|
||||||
|
"age" => [36, 63]
|
||||||
|
)
|
||||||
|
.expect("Should be able to create a simple dataframe"),
|
||||||
|
)
|
||||||
|
.into_value(Span::test_data()),
|
||||||
|
),
|
||||||
|
},
|
||||||
|
]
|
||||||
|
}
|
||||||
|
|
||||||
|
fn run(
|
||||||
|
&self,
|
||||||
|
plugin: &Self::Plugin,
|
||||||
|
engine: &EngineInterface,
|
||||||
|
call: &EvaluatedCall,
|
||||||
|
input: PipelineData,
|
||||||
|
) -> Result<PipelineData, LabeledError> {
|
||||||
|
match PolarsPluginObject::try_from_pipeline(plugin, input, call.head)? {
|
||||||
|
PolarsPluginObject::NuDataFrame(df) => command_eager(plugin, engine, call, df),
|
||||||
|
PolarsPluginObject::NuLazyFrame(lazy) => command_lazy(plugin, engine, call, lazy),
|
||||||
|
_ => Err(ShellError::GenericError {
|
||||||
|
error: "Must be a dataframe or lazy dataframe".into(),
|
||||||
|
msg: "".into(),
|
||||||
|
span: Some(call.head),
|
||||||
|
help: None,
|
||||||
|
inner: vec![],
|
||||||
|
}),
|
||||||
|
}
|
||||||
|
.map_err(LabeledError::from)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn command_eager(
|
||||||
|
plugin: &PolarsPlugin,
|
||||||
|
engine: &EngineInterface,
|
||||||
|
call: &EvaluatedCall,
|
||||||
|
df: NuDataFrame,
|
||||||
|
) -> Result<PipelineData, ShellError> {
|
||||||
|
let cols = call.rest::<String>(0)?;
|
||||||
|
let polars = df.to_polars();
|
||||||
|
let result: NuDataFrame = polars
|
||||||
|
.unnest(cols)
|
||||||
|
.map_err(|e| ShellError::GenericError {
|
||||||
|
error: format!("Error unnesting dataframe: {e}"),
|
||||||
|
msg: "".into(),
|
||||||
|
span: Some(call.head),
|
||||||
|
help: None,
|
||||||
|
inner: vec![],
|
||||||
|
})?
|
||||||
|
.into();
|
||||||
|
result.to_pipeline_data(plugin, engine, call.head)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn command_lazy(
|
||||||
|
plugin: &PolarsPlugin,
|
||||||
|
engine: &EngineInterface,
|
||||||
|
call: &EvaluatedCall,
|
||||||
|
df: NuLazyFrame,
|
||||||
|
) -> Result<PipelineData, ShellError> {
|
||||||
|
let cols = call.rest::<String>(0)?;
|
||||||
|
|
||||||
|
let polars = df.to_polars();
|
||||||
|
let result: NuLazyFrame = polars.unnest(cols).into();
|
||||||
|
result.to_pipeline_data(plugin, engine, call.head)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod test {
|
||||||
|
use crate::test::test_polars_plugin_command;
|
||||||
|
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_examples() -> Result<(), ShellError> {
|
||||||
|
test_polars_plugin_command(&UnnestDF)
|
||||||
|
}
|
||||||
|
}
|
Loading…
Reference in a new issue