extern/nushell
1
0
Fork 1
mirror of https://github.com/nushell/nushell.git synced 2024-11-07 17:14:23 +01:00
Code Issues Packages Projects Releases Wiki Activity

Fixing NuLazyFrame/NuDataFrame conversion issues (#12538)

Browse Source 
# Description

@maxim-uvarov brought up another case where converting back and forth
between eager and lazy dataframes was not working correctly:

```
> [[a b]; [6 2] [1 4] [4 1]] | polars into-lazy | polars append -c ([[a b]; [6 2] [1 4] [4 1]] | polars into-df)
Error: nu:🐚:cant_convert

  × Can't convert to NuDataFrame.
   ╭─[entry #1:1:49]
 1 │ [[a b]; [6 2] [1 4] [4 1]] | polars into-lazy | polars append -c ([[a b]; [6 2] [1 4] [4 1]] | polars into-df)
   ·                                                 ──────┬──────
   ·                                                       ╰── can't convert NuLazyFrameCustomValue to NuDataFrame
   ╰────
```

This pull request fixes this case and glaringly obvious similar cases I
could find.

Co-authored-by: Jack Wright <jack.wright@disqo.com>
This commit is contained in:
Jack Wright 2024-04-16 09:16:37 -07:00 committed by GitHub
parent 48e4448e55
commit a7a5ec31be
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
10 changed files with 25 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -126,7 +126,7 @@ fn command(
};
let df_other = NuDataFrame::try_from_value_coerce(plugin, &other, call.head)?;
let df = NuDataFrame::try_from_pipeline(plugin, input, call.head)?;
let df = NuDataFrame::try_from_pipeline_coerce(plugin, input, call.head)?;
let df = df.append_df(&df_other, axis, call.head)?;
df.to_pipeline_data(plugin, engine, call.head)

6
crates/nu_plugin_polars/src/dataframe/eager/first.rs
View File

@ -1,5 +1,5 @@
use crate::{
values::{Column, CustomValueSupport},
values::{Column, CustomValueSupport, NuLazyFrame},
PolarsPlugin,
};
@ -98,8 +98,8 @@ impl PluginCommand for FirstDF {
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let value = input.into_value(call.head);
if NuDataFrame::can_downcast(&value) {
let df = NuDataFrame::try_from_value(plugin, &value)?;
if NuDataFrame::can_downcast(&value) || NuLazyFrame::can_downcast(&value) {
let df = NuDataFrame::try_from_value_coerce(plugin, &value, call.head)?;
command(plugin, engine, call, df).map_err(|e| e.into())
} else {
let expr = NuExpression::try_from_value(plugin, &value)?;

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

@ -72,7 +72,7 @@ fn command(
let columns: Vec<Value> = call.rest(0)?;
let (col_string, col_span) = convert_columns_string(columns, call.head)?;
let df = NuDataFrame::try_from_pipeline(plugin, input, call.head)?;
let df = NuDataFrame::try_from_pipeline_coerce(plugin, input, call.head)?;
let df = df
.as_ref()

6
crates/nu_plugin_polars/src/dataframe/eager/last.rs
View File

@ -1,5 +1,5 @@
use crate::{
values::{Column, CustomValueSupport},
values::{Column, CustomValueSupport, NuLazyFrame},
PolarsPlugin,
};
@ -73,8 +73,8 @@ impl PluginCommand for LastDF {
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let value = input.into_value(call.head);
if NuDataFrame::can_downcast(&value) {
let df = NuDataFrame::try_from_value(plugin, &value)?;
if NuDataFrame::can_downcast(&value) || NuLazyFrame::can_downcast(&value) {
let df = NuDataFrame::try_from_value_coerce(plugin, &value, call.head)?;
command(plugin, engine, call, df).map_err(|e| e.into())
} else {
let expr = NuExpression::try_from_value(plugin, &value)?;

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

@ -141,7 +141,7 @@ fn command(
let (id_col_string, id_col_span) = convert_columns_string(id_col, call.head)?;
let (val_col_string, val_col_span) = convert_columns_string(val_col, call.head)?;
let df = NuDataFrame::try_from_pipeline(plugin, input, call.head)?;
let df = NuDataFrame::try_from_pipeline_coerce(plugin, input, call.head)?;
check_column_datatypes(df.as_ref(), &id_col_string, id_col_span)?;
check_column_datatypes(df.as_ref(), &val_col_string, val_col_span)?;

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

@ -92,7 +92,7 @@ fn command(
let replace: bool = call.has_flag("replace")?;
let shuffle: bool = call.has_flag("shuffle")?;
let df = NuDataFrame::try_from_pipeline(plugin, input, call.head)?;
let df = NuDataFrame::try_from_pipeline_coerce(plugin, input, call.head)?;
let df = match (rows, fraction) {
(Some(rows), None) => df

3
crates/nu_plugin_polars/src/dataframe/eager/take.rs
View File

@ -101,7 +101,8 @@ fn command(
) -> Result<PipelineData, ShellError> {
let index_value: Value = call.req(0)?;
let index_span = index_value.span();
let index = NuDataFrame::try_from_value(plugin, &index_value)?.as_series(index_span)?;
let index = NuDataFrame::try_from_value_coerce(plugin, &index_value, call.head)?
.as_series(index_span)?;
let casted = match index.dtype() {
DataType::UInt32 | DataType::UInt64 | DataType::Int32 | DataType::Int64 => index

10
crates/nu_plugin_polars/src/dataframe/eager/to_nu.rs
View File

@ -4,7 +4,11 @@ use nu_protocol::{
SyntaxShape, Type, Value,
};
use crate::{dataframe::values::NuExpression, values::CustomValueSupport, PolarsPlugin};
use crate::{
dataframe::values::NuExpression,
values::{CustomValueSupport, NuLazyFrame},
PolarsPlugin,
};
use super::super::values::NuDataFrame;
@ -86,7 +90,7 @@ impl PluginCommand for ToNu {
input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let value = input.into_value(call.head);
if NuDataFrame::can_downcast(&value) {
if NuDataFrame::can_downcast(&value) || NuLazyFrame::can_downcast(&value) {
dataframe_command(plugin, call, value)
} else {
expression_command(plugin, call, value)
@ -103,7 +107,7 @@ fn dataframe_command(
let rows: Option<usize> = call.get_flag("rows")?;
let tail: bool = call.has_flag("tail")?;
let df = NuDataFrame::try_from_value(plugin, &input)?;
let df = NuDataFrame::try_from_value_coerce(plugin, &input, call.head)?;
let values = if tail {
df.tail(rows, call.head)?

3
crates/nu_plugin_polars/src/dataframe/eager/with_column.rs
View File

@ -145,7 +145,8 @@ fn command_eager(
let df = lazy.collect(call.head)?;
df.to_pipeline_data(plugin, engine, call.head)
} else {
let mut other = NuDataFrame::try_from_value(plugin, &new_column)?.as_series(column_span)?;
let mut other = NuDataFrame::try_from_value_coerce(plugin, &new_column, call.head)?
.as_series(column_span)?;
let name = match call.get_flag::<String>("name")? {
Some(name) => name,

7
crates/nu_plugin_polars/src/dataframe/expressions/is_in.rs
View File

@ -165,7 +165,7 @@ fn command_df(
) -> Result<PipelineData, ShellError> {
let other_value: Value = call.req(0)?;
let other_span = other_value.span();
let other_df = NuDataFrame::try_from_value(plugin, &other_value)?;
let other_df = NuDataFrame::try_from_value_coerce(plugin, &other_value, call.head)?;
let other = other_df.as_series(other_span)?;
let series = df.as_series(call.head)?;
@ -181,8 +181,9 @@ fn command_df(
res.rename("is_in");
let df = NuDataFrame::try_from_series_vec(vec![res.into_series()], call.head)?;
df.to_pipeline_data(plugin, engine, call.head)
let mut new_df = NuDataFrame::try_from_series_vec(vec![res.into_series()], call.head)?;
new_df.from_lazy = df.from_lazy;
new_df.to_pipeline_data(plugin, engine, call.head)
}
#[cfg(test)]