diff --git a/crates/nu_plugin_polars/src/dataframe/eager/append.rs b/crates/nu_plugin_polars/src/dataframe/eager/append.rs index 8adc2c73db139..48c435baa36a1 100644 --- a/crates/nu_plugin_polars/src/dataframe/eager/append.rs +++ b/crates/nu_plugin_polars/src/dataframe/eager/append.rs @@ -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) diff --git a/crates/nu_plugin_polars/src/dataframe/eager/first.rs b/crates/nu_plugin_polars/src/dataframe/eager/first.rs index a05107a1b5b01..6bbdc66f75ed3 100644 --- a/crates/nu_plugin_polars/src/dataframe/eager/first.rs +++ b/crates/nu_plugin_polars/src/dataframe/eager/first.rs @@ -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 { 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)?; diff --git a/crates/nu_plugin_polars/src/dataframe/eager/get.rs b/crates/nu_plugin_polars/src/dataframe/eager/get.rs index 4f87c8acfd81b..34ba98154faf1 100644 --- a/crates/nu_plugin_polars/src/dataframe/eager/get.rs +++ b/crates/nu_plugin_polars/src/dataframe/eager/get.rs @@ -72,7 +72,7 @@ fn command( let columns: Vec = 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() diff --git a/crates/nu_plugin_polars/src/dataframe/eager/last.rs b/crates/nu_plugin_polars/src/dataframe/eager/last.rs index a9855222dea9d..840fe063e9f6b 100644 --- a/crates/nu_plugin_polars/src/dataframe/eager/last.rs +++ b/crates/nu_plugin_polars/src/dataframe/eager/last.rs @@ -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 { 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)?; diff --git a/crates/nu_plugin_polars/src/dataframe/eager/melt.rs b/crates/nu_plugin_polars/src/dataframe/eager/melt.rs index f01a5dfc3ff82..b69389ed241bb 100644 --- a/crates/nu_plugin_polars/src/dataframe/eager/melt.rs +++ b/crates/nu_plugin_polars/src/dataframe/eager/melt.rs @@ -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)?; diff --git a/crates/nu_plugin_polars/src/dataframe/eager/sample.rs b/crates/nu_plugin_polars/src/dataframe/eager/sample.rs index f46fa6582e077..48ca05959ac53 100644 --- a/crates/nu_plugin_polars/src/dataframe/eager/sample.rs +++ b/crates/nu_plugin_polars/src/dataframe/eager/sample.rs @@ -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 diff --git a/crates/nu_plugin_polars/src/dataframe/eager/take.rs b/crates/nu_plugin_polars/src/dataframe/eager/take.rs index b8a29198bead8..28b22095a1006 100644 --- a/crates/nu_plugin_polars/src/dataframe/eager/take.rs +++ b/crates/nu_plugin_polars/src/dataframe/eager/take.rs @@ -101,7 +101,8 @@ fn command( ) -> Result { 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 diff --git a/crates/nu_plugin_polars/src/dataframe/eager/to_nu.rs b/crates/nu_plugin_polars/src/dataframe/eager/to_nu.rs index 832396b139bde..55eb0d72055cb 100644 --- a/crates/nu_plugin_polars/src/dataframe/eager/to_nu.rs +++ b/crates/nu_plugin_polars/src/dataframe/eager/to_nu.rs @@ -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 { 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 = 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)? diff --git a/crates/nu_plugin_polars/src/dataframe/eager/with_column.rs b/crates/nu_plugin_polars/src/dataframe/eager/with_column.rs index 1ff0968b60e58..6973389729871 100644 --- a/crates/nu_plugin_polars/src/dataframe/eager/with_column.rs +++ b/crates/nu_plugin_polars/src/dataframe/eager/with_column.rs @@ -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::("name")? { Some(name) => name, diff --git a/crates/nu_plugin_polars/src/dataframe/expressions/is_in.rs b/crates/nu_plugin_polars/src/dataframe/expressions/is_in.rs index 1ecb7856c2240..b931a17764363 100644 --- a/crates/nu_plugin_polars/src/dataframe/expressions/is_in.rs +++ b/crates/nu_plugin_polars/src/dataframe/expressions/is_in.rs @@ -165,7 +165,7 @@ fn command_df( ) -> Result { 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)]