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

Incorporating API changes in polars 0.41

Browse Source
This commit is contained in:
Jack Wright 2024-06-27 11:03:19 -07:00
parent 7dd452cf46
commit f1e5bcd2c2
8 changed files with 86 additions and 33 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
crates/nu_plugin_polars/src/dataframe/eager/mod.rs
View File

@ -9,7 +9,6 @@ mod filter_with;
mod first;
mod get;
mod last;
mod unpivot;
mod open;
mod query_df;
mod rename;
@ -28,6 +27,7 @@ mod to_df;
mod to_json_lines;
mod to_nu;
mod to_parquet;
mod unpivot;
mod with_column;
use crate::PolarsPlugin;
@ -44,7 +44,6 @@ pub use filter_with::FilterWith;
pub use first::FirstDF;
pub use get::GetDF;
pub use last::LastDF;
pub use unpivot::UnpivotDF;
use nu_plugin::PluginCommand;
pub use query_df::QueryDf;
pub use rename::RenameDF;
@ -62,6 +61,7 @@ pub use to_df::ToDataFrame;
pub use to_json_lines::ToJsonLines;
pub use to_nu::ToNu;
pub use to_parquet::ToParquet;
pub use unpivot::UnpivotDF;
pub use with_column::WithColumn;
pub(crate) fn eager_commands() -> Vec<Box<dyn PluginCommand<Plugin = PolarsPlugin>>> {

2
crates/nu_plugin_polars/src/dataframe/eager/schema.rs
View File

@ -70,7 +70,7 @@ fn command(
let value: Value = schema.into();
Ok(PipelineData::Value(value, None))
}
PolarsPluginObject::NuLazyFrame(lazy) => {
PolarsPluginObject::NuLazyFrame(mut lazy) => {
let schema = lazy.schema()?;
let value: Value = schema.into();
Ok(PipelineData::Value(value, None))

3
crates/nu_plugin_polars/src/dataframe/lazy/aggregate.rs
View File

@ -196,7 +196,8 @@ fn get_col_name(expr: &Expr) -> Option<String> {
| Expr::Nth(_)
| Expr::SubPlan(_, _)
| Expr::IndexColumn(_)
| Expr::Selector(_) => None,
| Expr::Selector(_)
| Expr::Field(_) => None,
}
}

4
crates/nu_plugin_polars/src/dataframe/lazy/groupby.rs
View File

@ -148,11 +148,11 @@ fn command(
plugin: &PolarsPlugin,
engine: &EngineInterface,
call: &EvaluatedCall,
lazy: NuLazyFrame,
mut lazy: NuLazyFrame,
expressions: Vec<Expr>,
) -> Result<PipelineData, ShellError> {
let group_by = lazy.to_polars().group_by(expressions);
let group_by = NuLazyGroupBy::new(group_by, lazy.from_eager, lazy.schema()?);
let group_by = NuLazyGroupBy::new(group_by, lazy.from_eager, lazy.schema().clone()?);
group_by.to_pipeline_data(plugin, engine, call.head)
}

5
crates/nu_plugin_polars/src/dataframe/series/indexes/set_with_idx.rs
View File

@ -7,7 +7,10 @@ use nu_protocol::{
Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type,
Value,
};
use polars::{chunked_array::cast::CastOptions, prelude::{ChunkSet, DataType, IntoSeries}};
use polars::{
chunked_array::cast::CastOptions,
prelude::{ChunkSet, DataType, IntoSeries},
};
#[derive(Clone)]
pub struct SetWithIndex;

15
crates/nu_plugin_polars/src/dataframe/values/nu_dataframe/between_values.rs
View File

@ -65,14 +65,13 @@ pub(super) fn compute_between_series(
NuDataFrame::try_from_series(res, operation_span)
}
Operator::Math(Math::Multiply) => {
let mut res = (lhs * rhs)
.map_err(|e| ShellError::GenericError {
error: format!("Multiplication error: {e}"),
msg: "".into(),
span: Some(operation_span),
help: None,
inner: vec![],
})?;
let mut res = (lhs * rhs).map_err(|e| ShellError::GenericError {
error: format!("Multiplication error: {e}"),
msg: "".into(),
span: Some(operation_span),
help: None,
inner: vec![],
})?;
let name = format!("mul_{}_{}", lhs.name(), rhs.name());
res.rename(&name);
NuDataFrame::try_from_series(res, operation_span)

61
crates/nu_plugin_polars/src/dataframe/values/nu_expression/mod.rs
View File

@ -1,7 +1,10 @@
mod custom_value;
use nu_protocol::{record, ShellError, Span, Value};
use polars::prelude::{col, AggExpr, Expr, Literal};
use polars::{
chunked_array::cast::CastOptions,
prelude::{col, AggExpr, Expr, Literal},
};
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use uuid::Uuid;
@ -269,15 +272,23 @@ pub fn expr_to_value(expr: &Expr, span: Span) -> Result<Value, ShellError> {
Expr::Cast {
expr,
data_type,
strict,
} => Ok(Value::record(
record! {
"expr" => expr_to_value(expr.as_ref(), span)?,
"dtype" => Value::string(format!("{data_type:?}"), span),
"strict" => Value::bool(*strict, span),
},
span,
)),
options,
} => {
let cast_option_str = match options {
CastOptions::Strict => "STRICT",
CastOptions::NonStrict => "NON_STRICT",
CastOptions::Overflowing => "OVERFLOWING",
};
Ok(Value::record(
record! {
"expr" => expr_to_value(expr.as_ref(), span)?,
"dtype" => Value::string(format!("{data_type:?}"), span),
"cast_options" => Value::string(cast_option_str, span)
},
span,
))
}
Expr::Gather {
expr,
idx,
@ -388,6 +399,7 @@ pub fn expr_to_value(expr: &Expr, span: Span) -> Result<Value, ShellError> {
Expr::Window {
function,
partition_by,
order_by,
options,
} => {
let partition_by: Result<Vec<Value>, ShellError> = partition_by
@ -399,6 +411,23 @@ pub fn expr_to_value(expr: &Expr, span: Span) -> Result<Value, ShellError> {
record! {
"function" => expr_to_value(function, span)?,
"partition_by" => Value::list(partition_by?, span),
"order_by" => {
if let Some((order_expr, sort_options)) = order_by {
Value::record(record! {
"expr" => expr_to_value(order_expr.as_ref(), span)?,
"sort_options" => {
Value::record(record!(
"descending" => Value::bool(sort_options.descending, span),
"nulls_last"=> Value::bool(sort_options.nulls_last, span),
"multithreaded"=> Value::bool(sort_options.multithreaded, span),
"maintain_order"=> Value::bool(sort_options.maintain_order, span),
), span)
}
}, span)
} else {
Value::nothing(span)
}
},
"options" => Value::string(format!("{options:?}"), span),
},
span,
@ -424,6 +453,18 @@ pub fn expr_to_value(expr: &Expr, span: Span) -> Result<Value, ShellError> {
msg_span: span,
input_span: Span::unknown(),
}),
Expr::Field(column_name) => {
let fields: Vec<Value> = column_name
.iter()
.map(|s| Value::string(s.to_string(), span))
.collect();
Ok(Value::record(
record!(
"fields" => Value::list(fields, span)
),
span,
))
}
}
}

25
crates/nu_plugin_polars/src/dataframe/values/nu_lazyframe/mod.rs
View File

@ -77,14 +77,23 @@ impl NuLazyFrame {
Self::new(self.from_eager, new_frame)
}
pub fn schema(&self) -> Result<NuSchema, ShellError> {
let internal_schema = self.lazy.schema().map_err(|e| ShellError::GenericError {
error: "Error getting schema from lazy frame".into(),
msg: e.to_string(),
span: None,
help: None,
inner: vec![],
})?;
pub fn schema(&mut self) -> Result<NuSchema, ShellError> {
let internal_schema = Arc::get_mut(&mut self.lazy)
.ok_or(ShellError::GenericError {
error: "Error getting schema from lazy frame".into(),
msg: "LazyFrame is not mutable".into(),
span: None,
help: None,
inner: vec![],
})?
.schema()
.map_err(|e| ShellError::GenericError {
error: "Error getting schema from lazy frame".into(),
msg: e.to_string(),
span: None,
help: None,
inner: vec![],
})?;
Ok(internal_schema.into())
}