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:
Jack Wright 2024-10-16 19:24:14 -07:00 committed by sholderbach
parent e911ff4d67
commit dba2f6e0f8
2 changed files with 147 additions and 0 deletions

View File

@ -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),
]
}

View 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)
}
}