holm/nushell
1
1
Fork 0
forked from extern/nushell
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Upgrade to polars 0.43 (#14148)

Browse Source 
Upgrades the polars plugin to polars version 0.43
This commit is contained in:
Jack Wright
2024-10-23 10:14:24 -07:00
committed by GitHub
parent e7c4597ad0
commit ae54d05930
59 changed files with 530 additions and 398 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
crates/nu_plugin_polars/src/cache/get.rs vendored
View File

@ -42,7 +42,10 @@ impl PluginCommand for CacheGet {
polars store-ls | get key | first | polars store-get $in"#,
result: Some(
NuDataFrame::try_from_series_vec(
vec![Series::new("a", &[1_i64, 3]), Series::new("b", &[2_i64, 4])],
vec![
Series::new("a".into(), &[1_i64, 3]),
Series::new("b".into(), &[2_i64, 4]),
],
Span::test_data(),
)
.expect("could not create dataframe")

8
crates/nu_plugin_polars/src/dataframe/command/aggregation/cumulative.rs
View File

@ -82,9 +82,9 @@ impl PluginCommand for Cumulative {
vec![
Example {
description: "Cumulative sum for a column",
example: "[[a]; [1] [2] [3] [4] [5]]
| polars into-df
| polars select (polars col a | polars cumulative sum | polars as cum_a)
example: "[[a]; [1] [2] [3] [4] [5]]
| polars into-df
| polars select (polars col a | polars cumulative sum | polars as cum_a)
| polars collect",
result: Some(
NuDataFrame::try_from_columns(
@ -234,7 +234,7 @@ fn command_df(
})?;
let name = format!("{}_{}", series.name(), cum_type.to_str());
res.rename(&name);
res.rename(name.into());
let df = NuDataFrame::try_from_series_vec(vec![res.into_series()], call.head)?;
df.to_pipeline_data(plugin, engine, call.head)

2
crates/nu_plugin_polars/src/dataframe/command/aggregation/rolling.rs
View File

@ -169,7 +169,7 @@ fn command(
})?;
let name = format!("{}_{}", series.name(), roll_type.to_str());
res.rename(&name);
res.rename(name.into());
let df = NuDataFrame::try_from_series_vec(vec![res.into_series()], call.head)?;
df.to_pipeline_data(plugin, engine, call.head)

2
crates/nu_plugin_polars/src/dataframe/command/aggregation/value_counts.rs
View File

@ -92,7 +92,7 @@ fn command(
let series = df.as_series(call.head)?;
let res = series
.value_counts(sort, parallel, column, normalize)
.value_counts(sort, parallel, column.into(), normalize)
.map_err(|e| ShellError::GenericError {
error: "Error calculating value counts values".into(),
msg: e.to_string(),

2
crates/nu_plugin_polars/src/dataframe/command/boolean/arg_true.rs
View File

@ -84,7 +84,7 @@ fn command(
match columns.first() {
Some(column) => {
let expression = arg_where(col(column).eq(true)).alias("arg_true");
let expression = arg_where(col((*column).clone()).eq(true)).alias("arg_true");
let res: NuDataFrame = df
.as_ref()
.clone()

2
crates/nu_plugin_polars/src/dataframe/command/boolean/is_duplicated.rs
View File

@ -112,7 +112,7 @@ fn command(
})?
.into_series();
res.rename("is_duplicated");
res.rename("is_duplicated".into());
let df = NuDataFrame::try_from_series_vec(vec![res], call.head)?;
df.to_pipeline_data(plugin, engine, call.head)

2
crates/nu_plugin_polars/src/dataframe/command/boolean/is_in.rs
View File

@ -178,7 +178,7 @@ fn command_df(
})?
.into_series();
res.rename("is_in");
res.rename("is_in".into());
let mut new_df = NuDataFrame::try_from_series_vec(vec![res.into_series()], call.head)?;
new_df.from_lazy = df.from_lazy;

2
crates/nu_plugin_polars/src/dataframe/command/boolean/is_not_null.rs
View File

@ -108,7 +108,7 @@ fn command(
df: NuDataFrame,
) -> Result<PipelineData, ShellError> {
let mut res = df.as_series(call.head)?.is_not_null();
res.rename("is_not_null");
res.rename("is_not_null".into());
let df = NuDataFrame::try_from_series_vec(vec![res.into_series()], call.head)?;
df.to_pipeline_data(plugin, engine, call.head)

2
crates/nu_plugin_polars/src/dataframe/command/boolean/is_null.rs
View File

@ -110,7 +110,7 @@ fn command(
df: NuDataFrame,
) -> Result<PipelineData, ShellError> {
let mut res = df.as_series(call.head)?.is_null();
res.rename("is_null");
res.rename("is_null".into());
let df = NuDataFrame::try_from_series_vec(vec![res.into_series()], call.head)?;
df.to_pipeline_data(plugin, engine, call.head)

2
crates/nu_plugin_polars/src/dataframe/command/boolean/is_unique.rs
View File

@ -112,7 +112,7 @@ fn command(
})?
.into_series();
res.rename("is_unique");
res.rename("is_unique".into());
let df = NuDataFrame::try_from_series_vec(vec![res], call.head)?;
df.to_pipeline_data(plugin, engine, call.head)

2
crates/nu_plugin_polars/src/dataframe/command/boolean/set.rs
View File

@ -175,7 +175,7 @@ fn command(
})?;
let mut res = res.into_series();
res.rename("string");
res.rename("string".into());
NuDataFrame::try_from_series_vec(vec![res.into_series()], call.head)
}

2
crates/nu_plugin_polars/src/dataframe/command/core/columns.rs
View File

@ -59,7 +59,7 @@ fn command(
.as_ref()
.get_column_names()
.iter()
.map(|v| Value::string(*v, call.head))
.map(|v| Value::string(v.as_str(), call.head))
.collect();
let names = Value::list(names, call.head);

12
crates/nu_plugin_polars/src/dataframe/command/core/open.rs
View File

@ -24,7 +24,7 @@ use polars::{
lazy::frame::LazyJsonLineReader,
prelude::{
CsvEncoding, IpcReader, JsonFormat, JsonReader, LazyCsvReader, LazyFileListReader,
LazyFrame, ParquetReader, ScanArgsIpc, ScanArgsParquet, SerReader,
LazyFrame, ParquetReader, PlSmallStr, ScanArgsIpc, ScanArgsParquet, SerReader,
},
};
@ -190,6 +190,7 @@ fn from_parquet(
use_statistics: false,
hive_options: HiveOptions::default(),
glob: true,
include_file_paths: None,
};
let df: NuLazyFrame = LazyFrame::scan_parquet(file, args)
@ -286,8 +287,9 @@ fn from_arrow(
cache: true,
rechunk: false,
row_index: None,
memory_map: true,
cloud_options: None,
include_file_paths: None,
hive_options: HiveOptions::default(),
};
let df: NuLazyFrame = LazyFrame::scan_ipc(file, args)
@ -530,7 +532,11 @@ fn from_csv(
.with_infer_schema_length(Some(infer_schema))
.with_skip_rows(skip_rows.unwrap_or_default())
.with_schema(maybe_schema.map(|s| s.into()))
.with_columns(columns.map(|v| Arc::from(v.into_boxed_slice())))
.with_columns(
columns
.map(|v| v.iter().map(PlSmallStr::from).collect::<Vec<PlSmallStr>>())
.map(|v| Arc::from(v.into_boxed_slice())),
)
.map_parse_options(|options| {
options
.with_separator(

7
crates/nu_plugin_polars/src/dataframe/command/core/summary.rs
View File

@ -177,7 +177,8 @@ fn command(
let df = NuDataFrame::try_from_pipeline_coerce(plugin, input, call.head)?;
let names = ChunkedArray::<StringType>::from_slice_options("descriptor", &labels).into_series();
let names =
ChunkedArray::<StringType>::from_slice_options("descriptor".into(), &labels).into_series();
let head = std::iter::once(names);
@ -201,7 +202,7 @@ fn command(
.map(|q| {
col.quantile_reduce(q, QuantileInterpolOptions::default())
.ok()
.map(|s| s.into_series("quantile"))
.map(|s| s.into_series("quantile".into()))
.and_then(|ca| ca.cast(&DataType::Float64).ok())
.and_then(|ca| match ca.get(0) {
Ok(AnyValue::Float64(v)) => Some(v),
@ -217,7 +218,7 @@ fn command(
descriptors.push(max);
let name = format!("{} ({})", col.name(), col.dtype());
ChunkedArray::<Float64Type>::from_slice_options(&name, &descriptors).into_series()
ChunkedArray::<Float64Type>::from_slice_options(name.into(), &descriptors).into_series()
});
let res = head.chain(tail).collect::<Vec<Series>>();

18
crates/nu_plugin_polars/src/dataframe/command/core/to_df.rs
View File

@ -163,18 +163,18 @@ impl PluginCommand for ToDataFrame {
example: "[[a b c]; [1 {d: [1 2 3]} [10 11 12] ]]| polars into-df -s {a: u8, b: {d: list<u64>}, c: list<u8>}",
result: Some(
NuDataFrame::try_from_series_vec(vec![
Series::new("a", &[1u8]),
Series::new("a".into(), &[1u8]),
{
let dtype = DataType::Struct(vec![Field::new("a", DataType::List(Box::new(DataType::UInt64)))]);
let dtype = DataType::Struct(vec![Field::new("a".into(), DataType::List(Box::new(DataType::UInt64)))]);
let vals = vec![AnyValue::StructOwned(
Box::new((vec![AnyValue::List(Series::new("a", &[1u64, 2, 3]))], vec![Field::new("a", DataType::String)]))); 1];
Series::from_any_values_and_dtype("b", &vals, &dtype, false)
Box::new((vec![AnyValue::List(Series::new("a".into(), &[1u64, 2, 3]))], vec![Field::new("a".into(), DataType::String)]))); 1];
Series::from_any_values_and_dtype("b".into(), &vals, &dtype, false)
.expect("Struct series should not fail")
},
{
let dtype = DataType::List(Box::new(DataType::String));
let vals = vec![AnyValue::List(Series::new("c", &[10, 11, 12]))];
Series::from_any_values_and_dtype("c", &vals, &dtype, false)
let vals = vec![AnyValue::List(Series::new("c".into(), &[10, 11, 12]))];
Series::from_any_values_and_dtype("c".into(), &vals, &dtype, false)
.expect("List series should not fail")
}
], Span::test_data())
@ -186,9 +186,9 @@ impl PluginCommand for ToDataFrame {
description: "Convert to a dataframe and provide a schema that adds a new column",
example: r#"[[a b]; [1 "foo"] [2 "bar"]] | polars into-df -s {a: u8, b:str, c:i64} | polars fill-null 3"#,
result: Some(NuDataFrame::try_from_series_vec(vec![
Series::new("a", [1u8, 2]),
Series::new("b", ["foo", "bar"]),
Series::new("c", [3i64, 3]),
Series::new("a".into(), [1u8, 2]),
Series::new("b".into(), ["foo", "bar"]),
Series::new("c".into(), [3i64, 3]),
], Span::test_data())
.expect("simple df for test should not fail")
.into_value(Span::test_data()),

4
crates/nu_plugin_polars/src/dataframe/command/data/drop_duplicates.rs
View File

@ -49,7 +49,7 @@ impl PluginCommand for DropDuplicates {
fn examples(&self) -> Vec<Example> {
vec![Example {
description: "drop duplicates",
example: "[[a b]; [1 2] [3 4] [1 2]] | polars into-df
example: "[[a b]; [1 2] [3 4] [1 2]] | polars into-df
| polars drop-duplicates
| polars sort-by a",
result: Some(
@ -103,7 +103,7 @@ fn command(
let polars_df = df
.as_ref()
.unique(subset_slice, keep_strategy, None)
.unique_stable(subset_slice, keep_strategy, None)
.map_err(|e| ShellError::GenericError {
error: "Error dropping duplicates".into(),
msg: e.to_string(),

14
crates/nu_plugin_polars/src/dataframe/command/data/dummies.rs
View File

@ -38,10 +38,10 @@ impl PluginCommand for Dummies {
result: Some(
NuDataFrame::try_from_series_vec(
vec![
Series::new("a_1", &[1_u8, 0]),
Series::new("a_3", &[0_u8, 1]),
Series::new("b_2", &[1_u8, 0]),
Series::new("b_4", &[0_u8, 1]),
Series::new("a_1".into(), &[1_u8, 0]),
Series::new("a_3".into(), &[0_u8, 1]),
Series::new("b_2".into(), &[1_u8, 0]),
Series::new("b_4".into(), &[0_u8, 1]),
],
Span::test_data(),
)
@ -55,9 +55,9 @@ impl PluginCommand for Dummies {
result: Some(
NuDataFrame::try_from_series_vec(
vec![
Series::new("0_1", &[1_u8, 0, 0, 0, 0]),
Series::new("0_2", &[0_u8, 1, 1, 0, 0]),
Series::new("0_3", &[0_u8, 0, 0, 1, 1]),
Series::new("0_1".into(), &[1_u8, 0, 0, 0, 0]),
Series::new("0_2".into(), &[0_u8, 1, 1, 0, 0]),
Series::new("0_3".into(), &[0_u8, 0, 0, 1, 1]),
],
Span::test_data(),
)

3
crates/nu_plugin_polars/src/dataframe/command/data/mod.rs
View File

@ -69,7 +69,6 @@ pub use slice::SliceDF;
use sort_by_expr::LazySortBy;
pub use take::TakeDF;
pub use unique::Unique;
pub use unpivot::UnpivotDF;
pub use with_column::WithColumn;
pub(crate) fn data_commands() -> Vec<Box<dyn PluginCommand<Plugin = PolarsPlugin>>> {
@ -84,7 +83,7 @@ pub(crate) fn data_commands() -> Vec<Box<dyn PluginCommand<Plugin = PolarsPlugin
Box::new(filter_with::FilterWith),
Box::new(GetDF),
Box::new(pivot::PivotDF),
Box::new(UnpivotDF),
Box::new(unpivot::Unpivot),
Box::new(FirstDF),
Box::new(LastDF),
Box::new(len::ExprLen),

2
crates/nu_plugin_polars/src/dataframe/command/data/rename.rs
View File

@ -150,7 +150,7 @@ fn command_eager(
for (from, to) in columns.iter().zip(new_names.iter()) {
polars_df
.rename(from, to)
.rename(from, to.into())
.map_err(|e| ShellError::GenericError {
error: "Error renaming".into(),
msg: e.to_string(),

14
crates/nu_plugin_polars/src/dataframe/command/data/sample.rs
View File

@ -114,7 +114,12 @@ fn command(
let df = match (rows, fraction) {
(Some(rows), None) => df
.as_ref()
.sample_n(&Series::new("s", &[rows.item]), replace, shuffle, seed)
.sample_n(
&Series::new("s".into(), &[rows.item]),
replace,
shuffle,
seed,
)
.map_err(|e| ShellError::GenericError {
error: "Error creating sample".into(),
msg: e.to_string(),
@ -124,7 +129,12 @@ fn command(
}),
(None, Some(frac)) => df
.as_ref()
.sample_frac(&Series::new("frac", &[frac.item]), replace, shuffle, seed)
.sample_frac(
&Series::new("frac".into(), &[frac.item]),
replace,
shuffle,
seed,
)
.map_err(|e| ShellError::GenericError {
error: "Error creating sample".into(),
msg: e.to_string(),

7
crates/nu_plugin_polars/src/dataframe/command/data/sql_context.rs
View File

@ -98,11 +98,11 @@ impl SQLContext {
// Check for group by
// After projection since there might be number.
let group_by = match &select_stmt.group_by {
GroupByExpr::All =>
GroupByExpr::All(_) =>
Err(
PolarsError::ComputeError("Group-By Error: Only positive number or expression are supported, not all".into())
)?,
GroupByExpr::Expressions(expressions) => expressions
GroupByExpr::Expressions(expressions, _) => expressions
}
.iter()
.map(
@ -171,7 +171,8 @@ impl SQLContext {
.schema()
.get_at_index(shm_p)
.unwrap_or((&"".into(), &DataType::Null))
.0)
.0
.clone())
})
.collect::<Vec<_>>();
agg_df.select(final_proj)

18
crates/nu_plugin_polars/src/dataframe/command/data/unique.rs
View File

@ -1,5 +1,8 @@
use crate::{
dataframe::{utils::extract_strings, values::NuLazyFrame},
dataframe::{
utils::{extract_sm_strs, extract_strings},
values::NuLazyFrame,
},
values::{cant_convert_err, CustomValueSupport, PolarsPluginObject, PolarsPluginType},
PolarsPlugin,
};
@ -179,12 +182,7 @@ fn command_lazy(
) -> Result<PipelineData, ShellError> {
let last = call.has_flag("last")?;
let maintain = call.has_flag("maintain-order")?;
let subset: Option<Value> = call.get_flag("subset")?;
let subset = match subset {
Some(value) => Some(extract_strings(value)?),
None => None,
};
let strategy = if last {
UniqueKeepStrategy::Last
@ -194,8 +192,16 @@ fn command_lazy(
let lazy = lazy.to_polars();
let lazy: NuLazyFrame = if maintain {
let subset = match subset {
Some(value) => Some(extract_strings(value)?),
None => None,
};
lazy.unique(subset, strategy).into()
} else {
let subset = match subset {
Some(value) => Some(extract_sm_strs(value)?),
None => None,
};
lazy.unique_stable(subset, strategy).into()
};
lazy.to_pipeline_data(plugin, engine, call.head)

34
crates/nu_plugin_polars/src/dataframe/command/data/unpivot.rs
View File

@ -3,20 +3,21 @@ use nu_protocol::{
Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, Spanned,
SyntaxShape, Type, Value,
};
use polars::frame::explode::UnpivotArgs;
use polars::{frame::explode::UnpivotArgsIR, prelude::UnpivotArgsDSL};
use polars_ops::pivot::UnpivotDF;
use crate::{
dataframe::values::utils::convert_columns_string,
values::{CustomValueSupport, NuLazyFrame, PolarsPluginObject},
values::{utils::convert_columns_sm_str, CustomValueSupport, NuLazyFrame, PolarsPluginObject},
PolarsPlugin,
};
use crate::values::{Column, NuDataFrame};
#[derive(Clone)]
pub struct UnpivotDF;
pub struct Unpivot;
impl PluginCommand for UnpivotDF {
impl PluginCommand for Unpivot {
type Plugin = PolarsPlugin;
fn name(&self) -> &str {
@ -57,11 +58,6 @@ impl PluginCommand for UnpivotDF {
Type::Custom("dataframe".into()),
Type::Custom("dataframe".into()),
)
.switch(
"streamable",
"Whether or not to use the polars streaming engine. Only valid for lazy dataframes",
Some('t'),
)
.category(Category::Custom("dataframe".into()))
}
@ -220,14 +216,11 @@ fn command_eager(
check_column_datatypes(df.as_ref(), &index_col_string, index_col_span)?;
check_column_datatypes(df.as_ref(), &on_col_string, on_col_span)?;
let streamable = call.has_flag("streamable")?;
let args = UnpivotArgs {
let args = UnpivotArgsIR {
on: on_col_string.iter().map(Into::into).collect(),
index: index_col_string.iter().map(Into::into).collect(),
variable_name: variable_name.map(|s| s.item.into()),
value_name: value_name.map(|s| s.item.into()),
streamable,
};
let res = df
@ -254,20 +247,17 @@ fn command_lazy(
let index_col: Vec<Value> = call.get_flag("index")?.expect("required value");
let on_col: Vec<Value> = call.get_flag("on")?.expect("required value");
let (index_col_string, _index_col_span) = convert_columns_string(index_col, call.head)?;
let (on_col_string, _on_col_span) = convert_columns_string(on_col, call.head)?;
let (index_col_string, _index_col_span) = convert_columns_sm_str(index_col, call.head)?;
let (on_col_string, _on_col_span) = convert_columns_sm_str(on_col, call.head)?;
let value_name: Option<String> = call.get_flag("value-name")?;
let variable_name: Option<String> = call.get_flag("variable-name")?;
let streamable = call.has_flag("streamable")?;
let unpivot_args = UnpivotArgs {
on: on_col_string.iter().map(Into::into).collect(),
index: index_col_string.iter().map(Into::into).collect(),
let unpivot_args = UnpivotArgsDSL {
on: on_col_string.iter().cloned().map(Into::into).collect(),
index: index_col_string.iter().cloned().map(Into::into).collect(),
value_name: value_name.map(Into::into),
variable_name: variable_name.map(Into::into),
streamable,
};
let polars_df = df.to_polars().unpivot(unpivot_args);
@ -341,6 +331,6 @@ mod test {
#[test]
fn test_examples() -> Result<(), ShellError> {
test_polars_plugin_command(&UnpivotDF)
test_polars_plugin_command(&Unpivot)
}
}

2
crates/nu_plugin_polars/src/dataframe/command/data/with_column.rs
View File

@ -153,7 +153,7 @@ fn command_eager(
None => other.name().to_string(),
};
let series = other.rename(&name).clone();
let series = other.rename(name.into()).clone();
let mut polars_df = df.to_polars();
polars_df

2
crates/nu_plugin_polars/src/dataframe/command/datetime/as_date.rs
View File

@ -101,7 +101,7 @@ fn command(
})?
.into_series();
res.rename("date");
res.rename("date".into());
let df = NuDataFrame::try_from_series_vec(vec![res], call.head)?;
df.to_pipeline_data(plugin, engine, call.head)

2
crates/nu_plugin_polars/src/dataframe/command/datetime/as_datetime.rs
View File

@ -202,7 +202,7 @@ fn command(
})?
.into_series();
res.rename("datetime");
res.rename("datetime".into());
let df = NuDataFrame::try_from_series_vec(vec![res], call.head)?;
df.to_pipeline_data(plugin, engine, call.head)
}

16
crates/nu_plugin_polars/src/dataframe/command/datetime/datepart.rs
View File

@ -80,16 +80,16 @@ impl PluginCommand for ExprDatePart {
result: Some(
NuDataFrame::try_from_series_vec(
vec![
Series::new("datetime", &[dt.timestamp_nanos_opt()])
Series::new("datetime".into(), &[dt.timestamp_nanos_opt()])
.cast(&DataType::Datetime(TimeUnit::Nanoseconds, None))
.expect("Error casting to datetime type"),
Series::new("datetime_year", &[2021_i64]), // i32 was coerced to i64
Series::new("datetime_month", &[12_i8]),
Series::new("datetime_day", &[30_i8]),
Series::new("datetime_hour", &[1_i8]),
Series::new("datetime_minute", &[2_i8]),
Series::new("datetime_second", &[3_i8]),
Series::new("datetime_ns", &[123456789_i64]), // i32 was coerced to i64
Series::new("datetime_year".into(), &[2021_i64]), // i32 was coerced to i64
Series::new("datetime_month".into(), &[12_i8]),
Series::new("datetime_day".into(), &[30_i8]),
Series::new("datetime_hour".into(), &[1_i8]),
Series::new("datetime_minute".into(), &[2_i8]),
Series::new("datetime_second".into(), &[3_i8]),
Series::new("datetime_ns".into(), &[123456789_i64]), // i32 was coerced to i64
],
Span::test_data(),
)

2
crates/nu_plugin_polars/src/dataframe/command/datetime/get_day.rs
View File

@ -41,7 +41,7 @@ impl PluginCommand for GetDay {
let df = ([$dt $dt] | polars into-df);
$df | polars get-day"#,
result: Some(
NuDataFrame::try_from_series(Series::new("0", &[4i8, 4]), Span::test_data())
NuDataFrame::try_from_series(Series::new("0".into(), &[4i8, 4]), Span::test_data())
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),

9
crates/nu_plugin_polars/src/dataframe/command/datetime/get_hour.rs
View File

@ -41,9 +41,12 @@ impl PluginCommand for GetHour {
let df = ([$dt $dt] | polars into-df);
$df | polars get-hour"#,
result: Some(
NuDataFrame::try_from_series(Series::new("0", &[16i8, 16]), Span::test_data())
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
NuDataFrame::try_from_series(
Series::new("0".into(), &[16i8, 16]),
Span::test_data(),
)
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),
}]
}

9
crates/nu_plugin_polars/src/dataframe/command/datetime/get_minute.rs
View File

@ -39,9 +39,12 @@ impl PluginCommand for GetMinute {
let df = ([$dt $dt] | polars into-df);
$df | polars get-minute"#,
result: Some(
NuDataFrame::try_from_series(Series::new("0", &[39i8, 39]), Span::test_data())
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
NuDataFrame::try_from_series(
Series::new("0".into(), &[39i8, 39]),
Span::test_data(),
)
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),
}]
}

2
crates/nu_plugin_polars/src/dataframe/command/datetime/get_month.rs
View File

@ -41,7 +41,7 @@ impl PluginCommand for GetMonth {
let df = ([$dt $dt] | polars into-df);
$df | polars get-month"#,
result: Some(
NuDataFrame::try_from_series(Series::new("0", &[8i8, 8]), Span::test_data())
NuDataFrame::try_from_series(Series::new("0".into(), &[8i8, 8]), Span::test_data())
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),

9
crates/nu_plugin_polars/src/dataframe/command/datetime/get_nanosecond.rs
View File

@ -41,9 +41,12 @@ impl PluginCommand for GetNanosecond {
let df = ([$dt $dt] | polars into-df);
$df | polars get-nanosecond"#,
result: Some(
NuDataFrame::try_from_series(Series::new("0", &[0i32, 0]), Span::test_data())
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
NuDataFrame::try_from_series(
Series::new("0".into(), &[0i32, 0]),
Span::test_data(),
)
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),
}]
}

9
crates/nu_plugin_polars/src/dataframe/command/datetime/get_ordinal.rs
View File

@ -41,9 +41,12 @@ impl PluginCommand for GetOrdinal {
let df = ([$dt $dt] | polars into-df);
$df | polars get-ordinal"#,
result: Some(
NuDataFrame::try_from_series(Series::new("0", &[217i16, 217]), Span::test_data())
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
NuDataFrame::try_from_series(
Series::new("0".into(), &[217i16, 217]),
Span::test_data(),
)
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),
}]
}

9
crates/nu_plugin_polars/src/dataframe/command/datetime/get_second.rs
View File

@ -41,9 +41,12 @@ impl PluginCommand for GetSecond {
let df = ([$dt $dt] | polars into-df);
$df | polars get-second"#,
result: Some(
NuDataFrame::try_from_series(Series::new("0", &[18i8, 18]), Span::test_data())
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
NuDataFrame::try_from_series(
Series::new("0".into(), &[18i8, 18]),
Span::test_data(),
)
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),
}]
}

9
crates/nu_plugin_polars/src/dataframe/command/datetime/get_week.rs
View File

@ -41,9 +41,12 @@ impl PluginCommand for GetWeek {
let df = ([$dt $dt] | polars into-df);
$df | polars get-week"#,
result: Some(
NuDataFrame::try_from_series(Series::new("0", &[32i8, 32]), Span::test_data())
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
NuDataFrame::try_from_series(
Series::new("0".into(), &[32i8, 32]),
Span::test_data(),
)
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),
}]
}

2
crates/nu_plugin_polars/src/dataframe/command/datetime/get_weekday.rs
View File

@ -41,7 +41,7 @@ impl PluginCommand for GetWeekDay {
let df = ([$dt $dt] | polars into-df);
$df | polars get-weekday"#,
result: Some(
NuDataFrame::try_from_series(Series::new("0", &[2i8, 2]), Span::test_data())
NuDataFrame::try_from_series(Series::new("0".into(), &[2i8, 2]), Span::test_data())
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),

9
crates/nu_plugin_polars/src/dataframe/command/datetime/get_year.rs
View File

@ -41,9 +41,12 @@ impl PluginCommand for GetYear {
let df = ([$dt $dt] | polars into-df);
$df | polars get-year"#,
result: Some(
NuDataFrame::try_from_series(Series::new("0", &[2020i32, 2020]), Span::test_data())
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
NuDataFrame::try_from_series(
Series::new("0".into(), &[2020i32, 2020]),
Span::test_data(),
)
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),
}]
}

4
crates/nu_plugin_polars/src/dataframe/command/index/arg_max.rs
View File

@ -72,8 +72,8 @@ fn command(
let res = series.arg_max();
let chunked = match res {
Some(index) => UInt32Chunked::from_slice("arg_max", &[index as u32]),
None => UInt32Chunked::from_slice("arg_max", &[]),
Some(index) => UInt32Chunked::from_slice("arg_max".into(), &[index as u32]),
None => UInt32Chunked::from_slice("arg_max".into(), &[]),
};
let res = chunked.into_series();

4
crates/nu_plugin_polars/src/dataframe/command/index/arg_min.rs
View File

@ -72,8 +72,8 @@ fn command(
let res = series.arg_min();
let chunked = match res {
Some(index) => UInt32Chunked::from_slice("arg_min", &[index as u32]),
None => UInt32Chunked::from_slice("arg_min", &[]),
Some(index) => UInt32Chunked::from_slice("arg_min".into(), &[index as u32]),
None => UInt32Chunked::from_slice("arg_min".into(), &[]),
};
let res = chunked.into_series();

2
crates/nu_plugin_polars/src/dataframe/command/index/arg_sort.rs
View File

@ -119,7 +119,7 @@ fn command(
.as_series(call.head)?
.arg_sort(sort_options)
.into_series();
res.rename("arg_sort");
res.rename("arg_sort".into());
let df = NuDataFrame::try_from_series_vec(vec![res], call.head)?;
df.to_pipeline_data(plugin, engine, call.head)

2
crates/nu_plugin_polars/src/dataframe/command/index/arg_unique.rs
View File

@ -83,7 +83,7 @@ fn command(
inner: vec![],
})?
.into_series();
res.rename("arg_unique");
res.rename("arg_unique".into());
let df = NuDataFrame::try_from_series_vec(vec![res], call.head)?;
df.to_pipeline_data(plugin, engine, call.head)

2
crates/nu_plugin_polars/src/dataframe/command/index/set_with_idx.rs
View File

@ -195,7 +195,7 @@ fn command(
})?;
let mut res = res.into_series();
res.rename("string");
res.rename("string".into());
NuDataFrame::try_from_series_vec(vec![res.into_series()], call.head)
}

2
crates/nu_plugin_polars/src/dataframe/command/string/replace.rs
View File

@ -184,7 +184,7 @@ fn command_df(
inner: vec![],
})?;
res.rename(series.name());
res.rename(series.name().to_owned());
let df = NuDataFrame::try_from_series_vec(vec![res.into_series()], call.head)?;
df.to_pipeline_data(plugin, engine, call.head)

2
crates/nu_plugin_polars/src/dataframe/command/string/replace_all.rs
View File

@ -185,7 +185,7 @@ fn command_df(
inner: vec![],
})?;
res.rename(series.name());
res.rename(series.name().to_owned());
let df = NuDataFrame::try_from_series_vec(vec![res.into_series()], call.head)?;
df.to_pipeline_data(plugin, engine_state, call.head)

2
crates/nu_plugin_polars/src/dataframe/command/string/str_join.rs
View File

@ -159,7 +159,7 @@ fn command_df(
let mut res = chunked.concat(other_chunked);
res.rename(series.name());
res.rename(series.name().to_owned());
let df = NuDataFrame::try_from_series_vec(vec![res.into_series()], call.head)?;
df.to_pipeline_data(plugin, engine, call.head)

8
crates/nu_plugin_polars/src/dataframe/command/string/str_slice.rs
View File

@ -160,12 +160,12 @@ fn command_df(
df: NuDataFrame,
) -> Result<PipelineData, ShellError> {
let start: i64 = call.req(0)?;
let start = Series::new("", &[start]);
let start = Series::new("".into(), &[start]);
let length: Option<i64> = call.get_flag("length")?;
let length = match length {
Some(v) => Series::new("", &[v as u64]),
None => Series::new_null("", 1),
Some(v) => Series::new("".into(), &[v as u64]),
None => Series::new_null("".into(), 1),
};
let series = df.as_series(call.head)?;
@ -187,7 +187,7 @@ fn command_df(
help: None,
inner: vec![],
})?
.with_name(series.name());
.with_name(series.name().to_owned());
let df = NuDataFrame::try_from_series_vec(vec![res.into_series()], call.head)?;
df.to_pipeline_data(plugin, engine, call.head)

2
crates/nu_plugin_polars/src/dataframe/command/string/to_lowercase.rs
View File

@ -137,7 +137,7 @@ fn command_df(
})?;
let mut res = casted.to_lowercase();
res.rename(series.name());
res.rename(series.name().to_owned());
let df = NuDataFrame::try_from_series_vec(vec![res.into_series()], call.head)?;
df.to_pipeline_data(plugin, engine, call.head)

2
crates/nu_plugin_polars/src/dataframe/command/string/to_uppercase.rs
View File

@ -141,7 +141,7 @@ fn command_df(
})?;
let mut res = casted.to_uppercase();
res.rename(series.name());
res.rename(series.name().clone());
let df = NuDataFrame::try_from_series_vec(vec![res.into_series()], call.head)?;
df.to_pipeline_data(plugin, engine, call.head)

16
crates/nu_plugin_polars/src/dataframe/utils.rs
View File

@ -1,4 +1,5 @@
use nu_protocol::{FromValue, ShellError, Value};
use polars::prelude::PlSmallStr;
pub fn extract_strings(value: Value) -> Result<Vec<String>, ShellError> {
let span = value.span();
@ -14,3 +15,18 @@ pub fn extract_strings(value: Value) -> Result<Vec<String>, ShellError> {
}),
}
}
pub fn extract_sm_strs(value: Value) -> Result<Vec<PlSmallStr>, ShellError> {
let span = value.span();
match (
<String as FromValue>::from_value(value.clone()),
<Vec<String> as FromValue>::from_value(value),
) {
(Ok(col), Err(_)) => Ok(vec![col.into()]),
(Err(_), Ok(cols)) => Ok(cols.iter().map(PlSmallStr::from).collect()),
_ => Err(ShellError::IncompatibleParametersSingle {
msg: "Expected a string or list of strings".into(),
span,
}),
}
}

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

@ -49,7 +49,7 @@ pub(super) fn compute_between_series(
inner: vec![],
})?;
let name = format!("sum_{}_{}", lhs.name(), rhs.name());
res.rename(&name);
res.rename(name.into());
NuDataFrame::try_from_series(res, operation_span)
}
Operator::Math(Math::Minus) => {
@ -61,7 +61,7 @@ pub(super) fn compute_between_series(
inner: vec![],
})?;
let name = format!("sub_{}_{}", lhs.name(), rhs.name());
res.rename(&name);
res.rename(name.into());
NuDataFrame::try_from_series(res, operation_span)
}
Operator::Math(Math::Multiply) => {
@ -73,7 +73,7 @@ pub(super) fn compute_between_series(
inner: vec![],
})?;
let name = format!("mul_{}_{}", lhs.name(), rhs.name());
res.rename(&name);
res.rename(name.into());
NuDataFrame::try_from_series(res, operation_span)
}
Operator::Math(Math::Divide) => {
@ -81,7 +81,7 @@ pub(super) fn compute_between_series(
match res {
Ok(mut res) => {
let name = format!("div_{}_{}", lhs.name(), rhs.name());
res.rename(&name);
res.rename(name.into());
NuDataFrame::try_from_series(res, operation_span)
}
Err(e) => Err(ShellError::GenericError {
@ -132,7 +132,7 @@ pub(super) fn compute_between_series(
(Ok(l), Ok(r)) => {
let mut res = l.bitand(r).into_series();
let name = format!("and_{}_{}", lhs.name(), rhs.name());
res.rename(&name);
res.rename(name.into());
NuDataFrame::try_from_series(res, operation_span)
}
_ => Err(ShellError::GenericError {
@ -161,7 +161,7 @@ pub(super) fn compute_between_series(
(Ok(l), Ok(r)) => {
let mut res = l.bitor(r).into_series();
let name = format!("or_{}_{}", lhs.name(), rhs.name());
res.rename(&name);
res.rename(name.into());
NuDataFrame::try_from_series(res, operation_span)
}
_ => Err(ShellError::GenericError {
@ -211,7 +211,7 @@ where
})?
.into_series();
res.rename(name);
res.rename(name.into());
Ok(res)
}
@ -815,7 +815,7 @@ mod test {
#[test]
fn test_compute_between_series_comparisons() {
let series = Series::new("c", &[1, 2]);
let series = Series::new("c".into(), &[1, 2]);
let df = NuDataFrame::try_from_series_vec(vec![series], Span::test_data())
.expect("should be able to create a simple dataframe");
@ -848,7 +848,7 @@ mod test {
let result = result
.as_series(Span::test_data())
.expect("should be convert to a series");
assert_eq!(result, Series::new("neq_c_c", &[false, false]));
assert_eq!(result, Series::new("neq_c_c".into(), &[false, false]));
let op = Spanned {
item: Operator::Comparison(Comparison::Equal),
@ -859,7 +859,7 @@ mod test {
let result = result
.as_series(Span::test_data())
.expect("should be convert to a series");
assert_eq!(result, Series::new("eq_c_c", &[true, true]));
assert_eq!(result, Series::new("eq_c_c".into(), &[true, true]));
let op = Spanned {
item: Operator::Comparison(Comparison::LessThan),
@ -870,7 +870,7 @@ mod test {
let result = result
.as_series(Span::test_data())
.expect("should be convert to a series");
assert_eq!(result, Series::new("lt_c_c", &[false, false]));
assert_eq!(result, Series::new("lt_c_c".into(), &[false, false]));
let op = Spanned {
item: Operator::Comparison(Comparison::LessThanOrEqual),
@ -881,7 +881,7 @@ mod test {
let result = result
.as_series(Span::test_data())
.expect("should be convert to a series");
assert_eq!(result, Series::new("lte_c_c", &[true, true]));
assert_eq!(result, Series::new("lte_c_c".into(), &[true, true]));
let op = Spanned {
item: Operator::Comparison(Comparison::GreaterThan),
@ -892,7 +892,7 @@ mod test {
let result = result
.as_series(Span::test_data())
.expect("should be convert to a series");
assert_eq!(result, Series::new("gt_c_c", &[false, false]));
assert_eq!(result, Series::new("gt_c_c".into(), &[false, false]));
let op = Spanned {
item: Operator::Comparison(Comparison::GreaterThanOrEqual),
@ -903,6 +903,6 @@ mod test {
let result = result
.as_series(Span::test_data())
.expect("should be convert to a series");
assert_eq!(result, Series::new("gte_c_c", &[true, true]));
assert_eq!(result, Series::new("gte_c_c".into(), &[true, true]));
}
}

188
crates/nu_plugin_polars/src/dataframe/values/nu_dataframe/conversion.rs
View File

@ -7,14 +7,14 @@ use indexmap::map::{Entry, IndexMap};
use polars::chunked_array::builder::AnonymousOwnedListBuilder;
use polars::chunked_array::object::builder::ObjectChunkedBuilder;
use polars::chunked_array::ChunkedArray;
use polars::datatypes::AnyValue;
use polars::datatypes::{AnyValue, PlSmallStr};
use polars::export::arrow::Either;
use polars::prelude::{
DataFrame, DataType, DatetimeChunked, Float32Type, Float64Type, Int16Type, Int32Type,
Int64Type, Int8Type, IntoSeries, ListBooleanChunkedBuilder, ListBuilderTrait,
ChunkAnyValue, DataFrame, DataType, DatetimeChunked, Float32Type, Float64Type, Int16Type,
Int32Type, Int64Type, Int8Type, IntoSeries, ListBooleanChunkedBuilder, ListBuilderTrait,
ListPrimitiveChunkedBuilder, ListStringChunkedBuilder, ListType, NamedFrom, NewChunkedArray,
ObjectType, PolarsError, Schema, Series, StructChunked, TemporalMethods, TimeUnit, UInt16Type,
UInt32Type, UInt64Type, UInt8Type,
ObjectType, PolarsError, Schema, SchemaExt, Series, StructChunked, TemporalMethods, TimeUnit,
UInt16Type, UInt32Type, UInt64Type, UInt8Type,
};
use nu_protocol::{Record, ShellError, Span, Value};
@ -67,27 +67,27 @@ macro_rules! value_to_primitive {
#[derive(Debug)]
pub struct Column {
name: String,
name: PlSmallStr,
values: Vec<Value>,
}
impl Column {
pub fn new(name: impl Into<String>, values: Vec<Value>) -> Self {
pub fn new(name: impl Into<PlSmallStr>, values: Vec<Value>) -> Self {
Self {
name: name.into(),
values,
}
}
pub fn new_empty(name: String) -> Self {
pub fn new_empty(name: PlSmallStr) -> Self {
Self {
name,
values: Vec::new(),
}
}
pub fn name(&self) -> &str {
self.name.as_str()
pub fn name(&self) -> &PlSmallStr {
&self.name
}
}
@ -121,7 +121,7 @@ pub struct TypedColumn {
}
impl TypedColumn {
fn new_empty(name: String) -> Self {
fn new_empty(name: PlSmallStr) -> Self {
Self {
column: Column::new_empty(name),
column_type: None,
@ -143,7 +143,7 @@ impl DerefMut for TypedColumn {
}
}
pub type ColumnMap = IndexMap<String, TypedColumn>;
pub type ColumnMap = IndexMap<PlSmallStr, TypedColumn>;
pub fn create_column(
series: &Series,
@ -153,7 +153,7 @@ pub fn create_column(
) -> Result<Column, ShellError> {
let size = to_row - from_row;
let values = series_to_values(series, Some(from_row), Some(size), span)?;
Ok(Column::new(series.name(), values))
Ok(Column::new(series.name().clone(), values))
}
// Adds a separator to the vector of values using the column names from the
@ -168,7 +168,7 @@ pub fn add_separator(values: &mut Vec<Value>, df: &DataFrame, has_index: bool, s
for name in df.get_column_names() {
// there should only be one index field
record.push(name, Value::string("...", span))
record.push(name.as_str(), Value::string("...", span))
}
values.push(Value::record(record, span));
@ -181,7 +181,7 @@ pub fn insert_record(
maybe_schema: &Option<NuSchema>,
) -> Result<(), ShellError> {
for (col, value) in record {
insert_value(value, col, column_values, maybe_schema)?;
insert_value(value, col.into(), column_values, maybe_schema)?;
}
Ok(())
@ -189,7 +189,7 @@ pub fn insert_record(
pub fn insert_value(
value: Value,
key: String,
key: PlSmallStr,
column_values: &mut ColumnMap,
maybe_schema: &Option<NuSchema>,
) -> Result<(), ShellError> {
@ -201,7 +201,7 @@ pub fn insert_value(
// If we have a schema, use that for determining how things should be added to each column
if let Some(schema) = maybe_schema {
if let Some(field) = schema.schema.get_field(&key) {
col_val.column_type = Some(field.data_type().clone());
col_val.column_type = Some(field.dtype().clone());
col_val.values.push(value);
Ok(())
} else {
@ -219,7 +219,7 @@ pub fn insert_value(
else if col_val.values.is_empty() {
if let Some(schema) = maybe_schema {
if let Some(field) = schema.schema.get_field(&key) {
col_val.column_type = Some(field.data_type().clone());
col_val.column_type = Some(field.dtype().clone());
}
}
@ -284,7 +284,7 @@ fn value_to_data_type(value: &Value) -> DataType {
}
}
fn typed_column_to_series(name: &str, column: TypedColumn) -> Result<Series, ShellError> {
fn typed_column_to_series(name: PlSmallStr, column: TypedColumn) -> Result<Series, ShellError> {
if let Some(column_type) = &column.column_type {
match column_type {
DataType::Float32 => {
@ -426,13 +426,13 @@ fn typed_column_to_series(name: &str, column: TypedColumn) -> Result<Series, She
Ok(Series::new(name, series_values?))
}
DataType::List(list_type) => {
match input_type_list_to_series(name, list_type.as_ref(), &column.values) {
match input_type_list_to_series(&name, list_type.as_ref(), &column.values) {
Ok(series) => Ok(series),
Err(_) => {
// An error case will occur when there are lists of mixed types.
// If this happens, fallback to object list
input_type_list_to_series(
name,
&name,
&DataType::Object("unknown", None),
&column.values,
)
@ -489,7 +489,7 @@ fn typed_column_to_series(name: &str, column: TypedColumn) -> Result<Series, She
DataType::Struct(fields) => {
let schema = Some(NuSchema::new(Schema::from_iter(fields.clone())));
// let mut structs: Vec<Series> = Vec::new();
let mut structs: HashMap<String, Series> = HashMap::new();
let mut structs: HashMap<PlSmallStr, Series> = HashMap::new();
for v in column.values.iter() {
let mut column_values: ColumnMap = IndexMap::new();
@ -517,7 +517,7 @@ fn typed_column_to_series(name: &str, column: TypedColumn) -> Result<Series, She
inner: vec![],
})?;
} else {
structs.insert(name.to_string(), series.to_owned());
structs.insert(name.clone(), series.to_owned());
}
}
}
@ -525,15 +525,14 @@ fn typed_column_to_series(name: &str, column: TypedColumn) -> Result<Series, She
let structs: Vec<Series> = structs.into_values().collect();
let chunked =
StructChunked::new(column.name(), structs.as_slice()).map_err(|e| {
ShellError::GenericError {
StructChunked::from_series(column.name().to_owned(), structs.as_slice())
.map_err(|e| ShellError::GenericError {
error: format!("Error creating struct: {e}"),
msg: "".into(),
span: None,
help: None,
inner: vec![],
}
})?;
})?;
Ok(chunked.into_series())
}
_ => Err(ShellError::GenericError {
@ -561,7 +560,7 @@ fn typed_column_to_series(name: &str, column: TypedColumn) -> Result<Series, She
pub fn from_parsed_columns(column_values: ColumnMap) -> Result<NuDataFrame, ShellError> {
let mut df_series: Vec<Series> = Vec::new();
for (name, column) in column_values {
let series = typed_column_to_series(&name, column)?;
let series = typed_column_to_series(name, column)?;
df_series.push(series);
}
@ -576,7 +575,7 @@ pub fn from_parsed_columns(column_values: ColumnMap) -> Result<NuDataFrame, Shel
})
}
fn value_to_series(name: &str, values: &[Value]) -> Result<Series, ShellError> {
fn value_to_series(name: PlSmallStr, values: &[Value]) -> Result<Series, ShellError> {
let mut builder = ObjectChunkedBuilder::<DataFrameValue>::new(name, values.len());
for v in values {
@ -588,7 +587,7 @@ fn value_to_series(name: &str, values: &[Value]) -> Result<Series, ShellError> {
}
fn input_type_list_to_series(
name: &str,
name: &PlSmallStr,
data_type: &DataType,
values: &[Value],
) -> Result<Series, ShellError> {
@ -605,7 +604,7 @@ fn input_type_list_to_series(
macro_rules! primitive_list_series {
($list_type:ty, $vec_type:tt) => {{
let mut builder = ListPrimitiveChunkedBuilder::<$list_type>::new(
name,
name.clone(),
values.len(),
VALUES_CAPACITY,
data_type.clone(),
@ -628,7 +627,8 @@ fn input_type_list_to_series(
match *data_type {
// list of boolean values
DataType::Boolean => {
let mut builder = ListBooleanChunkedBuilder::new(name, values.len(), VALUES_CAPACITY);
let mut builder =
ListBooleanChunkedBuilder::new(name.clone(), values.len(), VALUES_CAPACITY);
for v in values {
let value_list = v
.as_list()?
@ -653,7 +653,8 @@ fn input_type_list_to_series(
DataType::Float32 => primitive_list_series!(Float32Type, f32),
DataType::Float64 => primitive_list_series!(Float64Type, f64),
DataType::String => {
let mut builder = ListStringChunkedBuilder::new(name, values.len(), VALUES_CAPACITY);
let mut builder =
ListStringChunkedBuilder::new(name.clone(), values.len(), VALUES_CAPACITY);
for v in values {
let value_list = v
.as_list()?
@ -668,7 +669,7 @@ fn input_type_list_to_series(
}
DataType::Date => {
let mut builder = AnonymousOwnedListBuilder::new(
name,
name.clone(),
values.len(),
Some(DataType::Datetime(TimeUnit::Nanoseconds, None)),
);
@ -682,7 +683,7 @@ fn input_type_list_to_series(
None
}
});
let dt_chunked = ChunkedArray::<Int64Type>::from_iter_options(&list_name, it)
let dt_chunked = ChunkedArray::<Int64Type>::from_iter_options(list_name.into(), it)
.into_datetime(TimeUnit::Nanoseconds, None);
builder
@ -702,7 +703,7 @@ fn input_type_list_to_series(
Ok(input_type_list_to_series(name, sub_list_type, values)?)
}
// treat everything else as an object
_ => Ok(value_to_series(name, values)?),
_ => Ok(value_to_series(name.clone(), values)?),
}
}
@ -1142,7 +1143,7 @@ fn series_to_values(
.collect::<Result<Vec<Value>, ShellError>>()?;
Ok(values)
}
DataType::Struct(polar_fields) => {
DataType::Struct(_) => {
let casted = series.struct_().map_err(|e| ShellError::GenericError {
error: "Error casting column to struct".into(),
msg: "".to_string(),
@ -1150,26 +1151,29 @@ fn series_to_values(
help: Some(e.to_string()),
inner: Vec::new(),
})?;
let it = casted.into_iter();
let values: Result<Vec<Value>, ShellError> =
if let (Some(size), Some(from_row)) = (maybe_size, maybe_from_row) {
Either::Left(it.skip(from_row).take(size))
} else {
Either::Right(it)
}
.map(|any_values| {
let record = polar_fields
.iter()
.zip(any_values)
.map(|(field, val)| {
any_value_to_value(val, span).map(|val| (field.name.to_string(), val))
})
.collect::<Result<_, _>>()?;
Ok(Value::record(record, span))
})
.collect();
values
let range = if let (Some(size), Some(from_row)) = (maybe_size, maybe_from_row) {
from_row..(from_row + size)
} else {
0..casted.len()
};
let mut values = Vec::with_capacity(casted.len());
for i in range {
let val = casted
.get_any_value(i)
.map_err(|e| ShellError::GenericError {
error: format!("Could not get struct value for index {i} - {e}"),
msg: "".into(),
span: None,
help: None,
inner: vec![],
})?;
values.push(any_value_to_value(&val, span)?)
}
Ok(values)
}
DataType::Time => {
let casted =
@ -1300,7 +1304,7 @@ fn nanos_from_timeunit(a: i64, time_unit: TimeUnit) -> i64 {
fn datetime_from_epoch_nanos(
nanos: i64,
timezone: &Option<String>,
timezone: &Option<PlSmallStr>,
span: Span,
) -> Result<DateTime<FixedOffset>, ShellError> {
let tz: Tz = if let Some(polars_tz) = timezone {
@ -1355,6 +1359,7 @@ where
mod tests {
use indexmap::indexmap;
use nu_protocol::record;
use polars::datatypes::CompatLevel;
use polars::export::arrow::array::{BooleanArray, PrimitiveArray};
use polars::prelude::Field;
use polars_io::prelude::StructArray;
@ -1374,7 +1379,7 @@ mod tests {
),
];
let column = Column {
name: "foo".to_string(),
name: "foo".into(),
values: values.clone(),
};
let typed_column = TypedColumn {
@ -1382,7 +1387,7 @@ mod tests {
column_type: Some(DataType::List(Box::new(DataType::String))),
};
let column_map = indexmap!("foo".to_string() => typed_column);
let column_map = indexmap!("foo".into() => typed_column);
let parsed_df = from_parsed_columns(column_map)?;
let parsed_columns = parsed_df.columns(Span::test_data())?;
assert_eq!(parsed_columns.len(), 1);
@ -1521,7 +1526,7 @@ mod tests {
Value::date(test_time, span)
);
let test_list_series = Series::new("int series", &[1, 2, 3]);
let test_list_series = Series::new("int series".into(), &[1, 2, 3]);
let comparison_list_series = Value::list(
vec![
Value::int(1, span),
@ -1541,8 +1546,8 @@ mod tests {
let field_name_0 = "num_field";
let field_name_1 = "bool_field";
let fields = vec![
Field::new(field_name_0, DataType::Int32),
Field::new(field_name_1, DataType::Boolean),
Field::new(field_name_0.into(), DataType::Int32),
Field::new(field_name_1.into(), DataType::Boolean),
];
let test_owned_struct = AnyValue::StructOwned(Box::new((values, fields.clone())));
let comparison_owned_record = Value::test_record(record!(
@ -1557,7 +1562,7 @@ mod tests {
let test_int_arr = PrimitiveArray::from([Some(1_i32)]);
let test_bool_arr = BooleanArray::from([Some(true)]);
let test_struct_arr = StructArray::new(
DataType::Struct(fields.clone()).to_arrow(true),
DataType::Struct(fields.clone()).to_arrow(CompatLevel::newest()),
vec![Box::new(test_int_arr), Box::new(test_bool_arr)],
None,
);
@ -1586,9 +1591,12 @@ mod tests {
column_type: Some(DataType::Float32),
};
let result = typed_column_to_series("foo", column)?;
let result = typed_column_to_series("foo".into(), column)?;
assert_eq!(result, Series::new("name", [Some(1.1f32), Some(2.0), None]));
assert_eq!(
result,
Series::new("name".into(), [Some(1.1f32), Some(2.0), None])
);
Ok(())
}
@ -1606,9 +1614,12 @@ mod tests {
column_type: Some(DataType::Float64),
};
let result = typed_column_to_series("foo", column)?;
let result = typed_column_to_series("foo".into(), column)?;
assert_eq!(result, Series::new("name", [Some(1.1f64), Some(2.0), None]));
assert_eq!(
result,
Series::new("name".into(), [Some(1.1f64), Some(2.0), None])
);
Ok(())
}
@ -1619,9 +1630,9 @@ mod tests {
column_type: Some(DataType::UInt8),
};
let result = typed_column_to_series("foo", column)?;
let result = typed_column_to_series("foo".into(), column)?;
assert_eq!(result, Series::new("name", [Some(1u8), None]));
assert_eq!(result, Series::new("name".into(), [Some(1u8), None]));
Ok(())
}
@ -1632,9 +1643,9 @@ mod tests {
column_type: Some(DataType::UInt16),
};
let result = typed_column_to_series("foo", column)?;
let result = typed_column_to_series("foo".into(), column)?;
assert_eq!(result, Series::new("name", [Some(1u16), None]));
assert_eq!(result, Series::new("name".into(), [Some(1u16), None]));
Ok(())
}
@ -1645,9 +1656,9 @@ mod tests {
column_type: Some(DataType::UInt32),
};
let result = typed_column_to_series("foo", column)?;
let result = typed_column_to_series("foo".into(), column)?;
assert_eq!(result, Series::new("name", [Some(1u32), None]));
assert_eq!(result, Series::new("name".into(), [Some(1u32), None]));
Ok(())
}
@ -1658,9 +1669,9 @@ mod tests {
column_type: Some(DataType::UInt64),
};
let result = typed_column_to_series("foo", column)?;
let result = typed_column_to_series("foo".into(), column)?;
assert_eq!(result, Series::new("name", [Some(1u64), None]));
assert_eq!(result, Series::new("name".into(), [Some(1u64), None]));
Ok(())
}
@ -1671,9 +1682,9 @@ mod tests {
column_type: Some(DataType::Int8),
};
let result = typed_column_to_series("foo", column)?;
let result = typed_column_to_series("foo".into(), column)?;
assert_eq!(result, Series::new("name", [Some(1i8), None]));
assert_eq!(result, Series::new("name".into(), [Some(1i8), None]));
Ok(())
}
@ -1684,9 +1695,9 @@ mod tests {
column_type: Some(DataType::Int16),
};
let result = typed_column_to_series("foo", column)?;
let result = typed_column_to_series("foo".into(), column)?;
assert_eq!(result, Series::new("name", [Some(1i16), None]));
assert_eq!(result, Series::new("name".into(), [Some(1i16), None]));
Ok(())
}
@ -1697,9 +1708,9 @@ mod tests {
column_type: Some(DataType::Int32),
};
let result = typed_column_to_series("foo", column)?;
let result = typed_column_to_series("foo".into(), column)?;
assert_eq!(result, Series::new("name", [Some(1i32), None]));
assert_eq!(result, Series::new("name".into(), [Some(1i32), None]));
Ok(())
}
@ -1710,9 +1721,9 @@ mod tests {
column_type: Some(DataType::Int64),
};
let result = typed_column_to_series("foo", column)?;
let result = typed_column_to_series("foo".into(), column)?;
assert_eq!(result, Series::new("name", [Some(1i64), None]));
assert_eq!(result, Series::new("name".into(), [Some(1i64), None]));
Ok(())
}
@ -1730,9 +1741,12 @@ mod tests {
column_type: Some(DataType::Boolean),
};
let result = typed_column_to_series("foo", column)?;
let result = typed_column_to_series("foo".into(), column)?;
assert_eq!(result, Series::new("name", [Some(true), Some(false), None]));
assert_eq!(
result,
Series::new("name".into(), [Some(true), Some(false), None])
);
Ok(())
}
@ -1746,11 +1760,11 @@ mod tests {
column_type: Some(DataType::String),
};
let result = typed_column_to_series("foo", column)?;
let result = typed_column_to_series("foo".into(), column)?;
assert_eq!(
result,
Series::new("name", [Some("barbaz".to_string()), None])
Series::new("name".into(), [Some("barbaz".to_string()), None])
);
Ok(())
}

7
crates/nu_plugin_polars/src/dataframe/values/nu_dataframe/mod.rs
View File

@ -181,7 +181,7 @@ impl NuDataFrame {
)?,
_ => {
let key = "0".to_string();
conversion::insert_value(value, key, &mut column_values, &maybe_schema)?
conversion::insert_value(value, key.into(), &mut column_values, &maybe_schema)?
}
}
}
@ -209,7 +209,7 @@ impl NuDataFrame {
let mut column_values: ColumnMap = IndexMap::new();
for column in columns {
let name = column.name().to_string();
let name = column.name().clone();
for value in column {
conversion::insert_value(value, name.clone(), &mut column_values, &maybe_schema)?;
}
@ -493,10 +493,9 @@ fn add_missing_columns(
})
.collect();
// todo - fix
let missing_exprs: Vec<Expr> = missing
.iter()
.map(|(name, dtype)| lit(Null {}).cast((*dtype).to_owned()).alias(name))
.map(|(name, dtype)| lit(Null {}).cast((*dtype).to_owned()).alias(*name))
.collect();
let df = if !missing.is_empty() {

4
crates/nu_plugin_polars/src/dataframe/values/nu_dataframe/operations.rs
View File

@ -126,7 +126,7 @@ impl NuDataFrame {
.iter()
.chain(other.df.get_columns())
.map(|s| {
let name = if columns.contains(&s.name()) {
let name = if columns.contains(&s.name().as_str()) {
format!("{}_{}", s.name(), "x")
} else {
columns.push(s.name());
@ -134,7 +134,7 @@ impl NuDataFrame {
};
let mut series = s.clone();
series.rename(&name);
series.rename(name.into());
series
})
.collect::<Vec<Series>>();

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

@ -146,7 +146,7 @@ pub fn expr_to_value(expr: &Expr, span: Span) -> Result<Value, ShellError> {
Expr::Alias(expr, alias) => Ok(Value::record(
record! {
"expr" => expr_to_value(expr.as_ref(), span)?,
"alias" => Value::string(alias.as_ref(), span),
"alias" => Value::string(alias.as_str(), span),
},
span,
)),
@ -271,7 +271,7 @@ pub fn expr_to_value(expr: &Expr, span: Span) -> Result<Value, ShellError> {
)),
Expr::Cast {
expr,
data_type,
dtype: data_type,
options,
} => {
let cast_option_str = match options {

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

@ -83,16 +83,15 @@ impl NuLazyFrame {
}
pub fn schema(&mut self) -> Result<NuSchema, ShellError> {
let internal_schema =
Arc::make_mut(&mut 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![],
})?;
let internal_schema = Arc::make_mut(&mut self.lazy)
.collect_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())
}

23
crates/nu_plugin_polars/src/dataframe/values/nu_schema.rs
View File

@ -3,7 +3,7 @@ use std::sync::Arc;
use nu_protocol::{ShellError, Span, Value};
use polars::{
datatypes::UnknownKind,
prelude::{DataType, Field, Schema, SchemaRef, TimeUnit},
prelude::{DataType, Field, PlSmallStr, Schema, SchemaExt, SchemaRef, TimeUnit},
};
#[derive(Debug, Clone)]
@ -49,7 +49,7 @@ fn fields_to_value(fields: impl Iterator<Item = Field>, span: Span) -> Value {
let record = fields
.map(|field| {
let col = field.name().to_string();
let val = dtype_to_value(field.data_type(), span);
let val = dtype_to_value(field.dtype(), span);
(col, val)
})
.collect();
@ -78,11 +78,11 @@ fn value_to_fields(value: &Value, span: Span) -> Result<Vec<Field>, ShellError>
Value::Record { .. } => {
let fields = value_to_fields(val, span)?;
let dtype = DataType::Struct(fields);
Ok(Field::new(col, dtype))
Ok(Field::new(col.into(), dtype))
}
_ => {
let dtype = str_to_dtype(&val.coerce_string()?, span)?;
Ok(Field::new(col, dtype))
Ok(Field::new(col.into(), dtype))
}
})
.collect::<Result<Vec<Field>, ShellError>>()?;
@ -150,7 +150,10 @@ pub fn str_to_dtype(dtype: &str, span: Span) -> Result<DataType, ShellError> {
} else {
Some(next.to_string())
};
Ok(DataType::Datetime(time_unit, timezone))
Ok(DataType::Datetime(
time_unit,
timezone.map(PlSmallStr::from),
))
}
_ if dtype.starts_with("duration") => {
let inner = dtype.trim_start_matches("duration<").trim_end_matches('>');
@ -215,13 +218,13 @@ mod test {
let schema = value_to_schema(&value, Span::unknown()).unwrap();
let expected = Schema::from_iter(vec![
Field::new("name", DataType::String),
Field::new("age", DataType::Int32),
Field::new("name".into(), DataType::String),
Field::new("age".into(), DataType::Int32),
Field::new(
"address",
"address".into(),
DataType::Struct(vec![
Field::new("street", DataType::String),
Field::new("city", DataType::String),
Field::new("street".into(), DataType::String),
Field::new("city".into(), DataType::String),
]),
),
]);

44
crates/nu_plugin_polars/src/dataframe/values/utils.rs
View File

@ -1,4 +1,5 @@
use nu_protocol::{ShellError, Span, Spanned, Value};
use polars::prelude::PlSmallStr;
// Default value used when selecting rows from dataframe
pub const DEFAULT_ROWS: usize = 5;
@ -88,3 +89,46 @@ pub(crate) fn convert_columns_string(
Ok((res, col_span))
}
// Converts a Vec<Value> to a Vec<String> with a Span marking the whole
// location of the columns for error referencing
// todo - fix
#[allow(dead_code)]
pub(crate) fn convert_columns_sm_str(
columns: Vec<Value>,
span: Span,
) -> Result<(Vec<PlSmallStr>, Span), ShellError> {
// First column span
let mut col_span = columns
.first()
.ok_or_else(|| ShellError::GenericError {
error: "Empty column list".into(),
msg: "Empty list found for command".into(),
span: Some(span),
help: None,
inner: vec![],
})
.map(|v| v.span())?;
let res = columns
.into_iter()
.map(|value| {
let span = value.span();
match value {
Value::String { val, .. } => {
col_span = col_span.merge(span);
Ok(val.into())
}
_ => Err(ShellError::GenericError {
error: "Incorrect column format".into(),
msg: "Only string as column name".into(),
span: Some(span),
help: None,
inner: vec![],
}),
}
})
.collect::<Result<Vec<PlSmallStr>, _>>()?;
Ok((res, col_span))
}