mirror of
https://github.com/nushell/nushell
synced 2024-12-26 04:53:09 +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_expr;
|
||||
mod take;
|
||||
mod unnest;
|
||||
mod unpivot;
|
||||
mod with_column;
|
||||
use filter::LazyFilter;
|
||||
|
@ -109,5 +110,6 @@ pub(crate) fn data_commands() -> Vec<Box<dyn PluginCommand<Plugin = PolarsPlugin
|
|||
Box::new(LazyFilter),
|
||||
Box::new(Shift),
|
||||
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