extern/nushell
1
0
Fork 1
mirror of https://github.com/nushell/nushell.git synced 2025-08-09 23:17:54 +02:00
Code Issues Packages Projects Releases Wiki Activity

Polars upgrade (#15852)

Browse Source 
# Description
Polars 0.48 upgrade

# User-Facing Changes
- (breaking change) Due to a change in behavior in polars, `polars
is-in` now only works as an expression.

---------

Co-authored-by: Jack Wright <jack.wright@nike.com>
This commit is contained in:
Jack Wright
2025-05-30 10:20:57 -07:00
committed by GitHub
parent 18ce5de500
commit d9ecb7da93
22 changed files with 382 additions and 271 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -186,7 +186,7 @@ fn get_col_name(expr: &Expr) -> Option<String> {
| Expr::Exclude(expr, _)
| Expr::Alias(expr, _)
| Expr::KeepName(expr)
| Expr::Explode(expr) => get_col_name(expr.as_ref()),
| Expr::Explode { input: expr, .. } => get_col_name(expr.as_ref()),
Expr::Ternary { .. }
| Expr::AnonymousFunction { .. }
| Expr::Function { .. }

12
crates/nu_plugin_polars/src/dataframe/command/aggregation/over.rs
View File

@ -5,7 +5,8 @@ use crate::{
};
use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand};
use nu_protocol::{
Category, Example, LabeledError, PipelineData, Signature, Span, SyntaxShape, Type, Value,
Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type,
Value,
};
use polars::df;
@ -94,7 +95,14 @@ impl PluginCommand for Over {
PolarsPluginObject::NuExpression(expr) => {
let expr: NuExpression = expr
.into_polars()
.over_with_options(expressions, None, Default::default())
.over_with_options(Some(expressions), None, Default::default())
.map_err(|e| ShellError::GenericError {
error: format!("Error applying over expression: {e}"),
msg: "".into(),
span: Some(call.head),
help: None,
inner: vec![],
})?
.into();
expr.to_pipeline_data(plugin, engine, call.head)
}

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

@ -8,7 +8,7 @@ use nu_protocol::{
Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type,
Value,
};
use polars::prelude::{DataType, IntoSeries, is_in, lit};
use polars::prelude::{DataType, lit};
#[derive(Clone)]
pub struct ExprIsIn;
@ -27,80 +27,48 @@ impl PluginCommand for ExprIsIn {
fn signature(&self) -> Signature {
Signature::build(self.name())
.required("list", SyntaxShape::Any, "List to check if values are in")
.input_output_types(vec![
(
Type::Custom("expression".into()),
Type::Custom("expression".into()),
),
(
Type::Custom("dataframe".into()),
Type::Custom("dataframe".into()),
),
])
.input_output_types(vec![(
Type::Custom("expression".into()),
Type::Custom("expression".into()),
)])
.category(Category::Custom("expression".into()))
}
fn examples(&self) -> Vec<Example> {
vec![
Example {
description: "Creates a is-in expression",
example: r#"let df = ([[a b]; [one 1] [two 2] [three 3]] | polars into-df);
vec![Example {
description: "Creates a is-in expression",
example: r#"let df = ([[a b]; [one 1] [two 2] [three 3]] | polars into-df);
$df | polars with-column (polars col a | polars is-in [one two] | polars as a_in)"#,
result: Some(
NuDataFrame::try_from_columns(
vec![
Column::new(
"a".to_string(),
vec![
Value::test_string("one"),
Value::test_string("two"),
Value::test_string("three"),
],
),
Column::new(
"b".to_string(),
vec![Value::test_int(1), Value::test_int(2), Value::test_int(3)],
),
Column::new(
"a_in".to_string(),
vec![
Value::test_bool(true),
Value::test_bool(true),
Value::test_bool(false),
],
),
],
None,
)
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),
},
Example {
description: "Checks if elements from a series are contained in right series",
example: r#"let other = ([1 3 6] | polars into-df);
[5 6 6 6 8 8 8] | polars into-df | polars is-in $other"#,
result: Some(
NuDataFrame::try_from_columns(
vec![Column::new(
"is_in".to_string(),
result: Some(
NuDataFrame::try_from_columns(
vec![
Column::new(
"a".to_string(),
vec![
Value::test_string("one"),
Value::test_string("two"),
Value::test_string("three"),
],
),
Column::new(
"b".to_string(),
vec![Value::test_int(1), Value::test_int(2), Value::test_int(3)],
),
Column::new(
"a_in".to_string(),
vec![
Value::test_bool(false),
Value::test_bool(true),
Value::test_bool(true),
Value::test_bool(true),
Value::test_bool(false),
Value::test_bool(false),
Value::test_bool(false),
],
)],
None,
)
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),
},
]
),
],
None,
)
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),
}]
}
fn search_terms(&self) -> Vec<&str> {
@ -117,10 +85,6 @@ impl PluginCommand for ExprIsIn {
let metadata = input.metadata();
let value = input.into_value(call.head)?;
match PolarsPluginObject::try_from_value(plugin, &value)? {
PolarsPluginObject::NuDataFrame(df) => command_df(plugin, engine, call, df),
PolarsPluginObject::NuLazyFrame(lazy) => {
command_df(plugin, engine, call, lazy.collect(call.head)?)
}
PolarsPluginObject::NuExpression(expr) => command_expr(plugin, engine, call, expr),
_ => Err(cant_convert_err(
&value,
@ -154,39 +118,11 @@ fn command_expr(
});
}
let expr: NuExpression = expr.into_polars().is_in(lit(list)).into();
// todo - at some point we should probably make this consistent with python api
let expr: NuExpression = expr.into_polars().is_in(lit(list).implode(), true).into();
expr.to_pipeline_data(plugin, engine, call.head)
}
fn command_df(
plugin: &PolarsPlugin,
engine: &EngineInterface,
call: &EvaluatedCall,
df: NuDataFrame,
) -> Result<PipelineData, ShellError> {
let other_value: Value = call.req(0)?;
let other_span = other_value.span();
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)?;
let mut res = is_in(&series, &other)
.map_err(|e| ShellError::GenericError {
error: "Error finding in other".into(),
msg: e.to_string(),
span: Some(call.head),
help: None,
inner: vec![],
})?
.into_series();
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;
new_df.to_pipeline_data(plugin, engine, call.head)
}
#[cfg(test)]
mod test {
use super::*;

13
crates/nu_plugin_polars/src/dataframe/command/core/save/arrow.rs
View File

@ -1,8 +1,9 @@
use std::fs::File;
use log::debug;
use nu_plugin::EvaluatedCall;
use nu_protocol::ShellError;
use polars::prelude::{IpcWriter, SerWriter};
use polars::prelude::{IpcWriter, SerWriter, SinkOptions};
use polars_io::ipc::IpcWriterOptions;
use crate::{
@ -10,21 +11,27 @@ use crate::{
values::{NuDataFrame, NuLazyFrame},
};
use super::polars_file_save_error;
use super::{polars_file_save_error, sink_target_from_string};
pub(crate) fn command_lazy(
_call: &EvaluatedCall,
lazy: &NuLazyFrame,
resource: Resource,
) -> Result<(), ShellError> {
let file_path = resource.path;
let file_path = sink_target_from_string(resource.path.clone());
let file_span = resource.span;
debug!("Writing ipc file {}", resource.path);
lazy.to_polars()
.sink_ipc(
file_path,
IpcWriterOptions::default(),
resource.cloud_options,
SinkOptions::default(),
)
.and_then(|l| l.collect())
.map(|_| {
debug!("Wrote ipc file {}", resource.path);
})
.map_err(|e| polars_file_save_error(e, file_span))
}

19
crates/nu_plugin_polars/src/dataframe/command/core/save/csv.rs
View File

@ -1,12 +1,13 @@
use std::fs::File;
use log::debug;
use nu_plugin::EvaluatedCall;
use nu_protocol::{ShellError, Spanned};
use polars::prelude::{CsvWriter, SerWriter};
use polars::prelude::{CsvWriter, SerWriter, SinkOptions};
use polars_io::csv::write::{CsvWriterOptions, SerializeOptions};
use crate::{
command::core::resource::Resource,
command::core::{resource::Resource, save::sink_target_from_string},
values::{NuDataFrame, NuLazyFrame},
};
@ -17,8 +18,9 @@ pub(crate) fn command_lazy(
lazy: &NuLazyFrame,
resource: Resource,
) -> Result<(), ShellError> {
let file_path = resource.path;
let file_path = sink_target_from_string(resource.path.clone());
let file_span = resource.span;
debug!("Writing csv file {}", resource.path);
let delimiter: Option<Spanned<String>> = call.get_flag("csv-delimiter")?;
let separator = delimiter
.and_then(|d| d.item.chars().next().map(|c| c as u8))
@ -36,8 +38,17 @@ pub(crate) fn command_lazy(
};
lazy.to_polars()
.sink_csv(file_path, options, resource.cloud_options)
.sink_csv(
file_path,
options,
resource.cloud_options,
SinkOptions::default(),
)
.and_then(|l| l.collect())
.map_err(|e| polars_file_save_error(e, file_span))
.map(|_| {
debug!("Wrote parquet file {}", resource.path);
})
}
pub(crate) fn command_eager(

11
crates/nu_plugin_polars/src/dataframe/command/core/save/mod.rs
View File

@ -4,7 +4,7 @@ mod csv;
mod ndjson;
mod parquet;
use std::path::PathBuf;
use std::{path::PathBuf, sync::Arc};
use crate::{
PolarsPlugin,
@ -19,7 +19,7 @@ use nu_protocol::{
Signature, Span, Spanned, SyntaxShape, Type,
shell_error::{self, io::IoError},
};
use polars::error::PolarsError;
use polars::{error::PolarsError, prelude::SinkTarget};
#[derive(Clone)]
pub struct SaveDF;
@ -272,6 +272,13 @@ pub fn unknown_file_save_error(span: Span) -> ShellError {
}
}
pub(crate) fn sink_target_from_string(path: String) -> SinkTarget {
let path = PathBuf::from(path);
let target = SinkTarget::Path(Arc::new(path));
debug!("Sink target: {target:?}");
target
}
#[cfg(test)]
pub(crate) mod test {
use nu_plugin_test_support::PluginTest;

17
crates/nu_plugin_polars/src/dataframe/command/core/save/ndjson.rs
View File

@ -1,12 +1,13 @@
use std::{fs::File, io::BufWriter};
use log::debug;
use nu_plugin::EvaluatedCall;
use nu_protocol::ShellError;
use polars::prelude::{JsonWriter, SerWriter};
use polars::prelude::{JsonWriter, SerWriter, SinkOptions};
use polars_io::json::JsonWriterOptions;
use crate::{
command::core::resource::Resource,
command::core::{resource::Resource, save::sink_target_from_string},
values::{NuDataFrame, NuLazyFrame},
};
@ -17,15 +18,21 @@ pub(crate) fn command_lazy(
lazy: &NuLazyFrame,
resource: Resource,
) -> Result<(), ShellError> {
let file_path = resource.path;
let file_path = sink_target_from_string(resource.path.clone());
let file_span = resource.span;
debug!("Writing ndjson file {}", resource.path);
lazy.to_polars()
.sink_json(
file_path,
JsonWriterOptions::default(),
resource.cloud_options,
SinkOptions::default(),
)
.and_then(|l| l.collect())
.map_err(|e| polars_file_save_error(e, file_span))
.map(|_| {
debug!("Wrote ndjson file {}", resource.path);
})
}
pub(crate) fn command_eager(df: &NuDataFrame, resource: Resource) -> Result<(), ShellError> {
@ -58,12 +65,12 @@ pub mod test {
use crate::command::core::save::test::{test_eager_save, test_lazy_save};
#[test]
pub fn test_arrow_eager_save() -> Result<(), Box<dyn std::error::Error>> {
pub fn test_ndjson_eager_save() -> Result<(), Box<dyn std::error::Error>> {
test_eager_save("ndjson")
}
#[test]
pub fn test_arrow_lazy_save() -> Result<(), Box<dyn std::error::Error>> {
pub fn test_ndjson_lazy_save() -> Result<(), Box<dyn std::error::Error>> {
test_lazy_save("ndjson")
}
}

16
crates/nu_plugin_polars/src/dataframe/command/core/save/parquet.rs
View File

@ -3,11 +3,11 @@ use std::fs::File;
use log::debug;
use nu_plugin::EvaluatedCall;
use nu_protocol::ShellError;
use polars::prelude::ParquetWriter;
use polars::prelude::{ParquetWriter, SinkOptions};
use polars_io::parquet::write::ParquetWriteOptions;
use crate::{
command::core::resource::Resource,
command::core::{resource::Resource, save::sink_target_from_string},
values::{NuDataFrame, NuLazyFrame},
};
@ -18,16 +18,22 @@ pub(crate) fn command_lazy(
lazy: &NuLazyFrame,
resource: Resource,
) -> Result<(), ShellError> {
let file_path = resource.path;
let file_path = sink_target_from_string(resource.path.clone());
let file_span = resource.span;
debug!("Writing parquet file {file_path}");
debug!("Writing parquet file {}", resource.path);
lazy.to_polars()
.sink_parquet(
&file_path,
file_path,
ParquetWriteOptions::default(),
resource.cloud_options,
SinkOptions::default(),
)
.and_then(|l| l.collect())
.map_err(|e| polars_file_save_error(e, file_span))
.map(|_| {
debug!("Wrote parquet file {}", resource.path);
})
}
pub(crate) fn command_eager(df: &NuDataFrame, resource: Resource) -> Result<(), ShellError> {

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

@ -188,7 +188,7 @@ fn command(
let tail = df
.as_ref()
.iter()
.filter(|col| !matches!(col.dtype(), &DataType::Object("object", _)))
.filter(|col| !matches!(col.dtype(), &DataType::Object("object")))
.map(|col| {
let count = col.len() as f64;

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

@ -102,7 +102,7 @@ fn literal_expr(value: &SqlValue) -> Result<Expr> {
SqlValue::HexStringLiteral(s) => lit(s.clone()),
SqlValue::DoubleQuotedString(s) => lit(s.clone()),
SqlValue::Boolean(b) => lit(*b),
SqlValue::Null => Expr::Literal(LiteralValue::Null),
SqlValue::Null => Expr::Literal(LiteralValue::untyped_null()),
_ => {
return Err(PolarsError::ComputeError(
format!("Parsing SQL Value {value:?} was not supported in polars-sql yet!").into(),

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

@ -6,6 +6,7 @@ use crate::{
},
};
use chrono::DateTime;
use polars_plan::plans::DynLiteralValue;
use std::sync::Arc;
use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand};
@ -294,7 +295,9 @@ fn command(
None,
None,
options,
Expr::Literal(LiteralValue::String(PlSmallStr::from_string(ambiguous))),
Expr::Literal(LiteralValue::Dyn(DynLiteralValue::Str(
PlSmallStr::from_string(ambiguous),
))),
)
.into();
res.to_pipeline_data(plugin, engine, call.head)
@ -324,7 +327,9 @@ fn command_lazy(
None,
None,
options,
Expr::Literal(LiteralValue::String(PlSmallStr::from_string(ambiguous))),
Expr::Literal(LiteralValue::Dyn(DynLiteralValue::Str(
PlSmallStr::from_string(ambiguous),
))),
)]),
)
.to_pipeline_data(plugin, engine, call.head)

22
crates/nu_plugin_polars/src/dataframe/command/datetime/convert_time_zone.rs
View File

@ -7,12 +7,15 @@ use crate::{
use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand};
use nu_protocol::{
Category, Example, LabeledError, PipelineData, Signature, Span, SyntaxShape, Type, Value,
Category, Example, LabeledError, PipelineData, Signature, Span, Spanned, SyntaxShape, Type,
Value,
};
use chrono::DateTime;
use polars::prelude::*;
use super::timezone_from_str;
#[derive(Clone)]
pub struct ConvertTimeZone;
@ -76,7 +79,8 @@ impl PluginCommand for ConvertTimeZone {
"datetime".into(),
DataType::Datetime(
TimeUnit::Nanoseconds,
Some(PlSmallStr::from_static("Europe/Lisbon")),
TimeZone::opt_try_new(Some("Europe/Lisbon"))
.expect("timezone should be valid"),
),
),
])))),
@ -118,7 +122,8 @@ impl PluginCommand for ConvertTimeZone {
"datetime".into(),
DataType::Datetime(
TimeUnit::Nanoseconds,
Some(PlSmallStr::from_static("America/New_York")),
TimeZone::opt_try_new(Some("America/New_York"))
.expect("timezone should be valid"),
),
),
])))),
@ -142,12 +147,11 @@ impl PluginCommand for ConvertTimeZone {
match PolarsPluginObject::try_from_value(plugin, &value)? {
PolarsPluginObject::NuExpression(expr) => {
let time_zone: String = call.req(0)?;
let expr: NuExpression = expr
.into_polars()
.dt()
.convert_time_zone(PlSmallStr::from_str(&time_zone))
.into();
let time_zone_spanned: Spanned<String> = call.req(0)?;
let time_zone =
timezone_from_str(&time_zone_spanned.item, Some(time_zone_spanned.span))?;
let expr: NuExpression =
expr.into_polars().dt().convert_time_zone(time_zone).into();
expr.to_pipeline_data(plugin, engine, call.head)
}
_ => Err(cant_convert_err(&value, &[PolarsPluginType::NuExpression])),

44
crates/nu_plugin_polars/src/dataframe/command/datetime/mod.rs
View File

@ -33,6 +33,8 @@ pub use get_second::GetSecond;
pub use get_week::GetWeek;
pub use get_weekday::GetWeekDay;
pub use get_year::GetYear;
use nu_protocol::{ShellError, Span};
use polars::prelude::{PlSmallStr, TimeZone};
pub use replace_time_zone::ReplaceTimeZone;
pub use strftime::StrFTime;
pub use truncate::Truncate;
@ -58,3 +60,45 @@ pub(crate) fn datetime_commands() -> Vec<Box<dyn PluginCommand<Plugin = PolarsPl
Box::new(Truncate),
]
}
pub fn timezone_from_str(zone_str: &str, span: Option<Span>) -> Result<TimeZone, ShellError> {
TimeZone::opt_try_new(Some(PlSmallStr::from_str(zone_str)))
.map_err(|e| ShellError::GenericError {
error: format!("Invalid timezone: {} : {}", zone_str, e),
msg: "".into(),
span,
help: None,
inner: vec![],
})?
.ok_or(ShellError::GenericError {
error: format!("Invalid timezone {}", zone_str),
msg: "".into(),
span,
help: None,
inner: vec![],
})
}
pub fn timezone_utc() -> TimeZone {
TimeZone::opt_try_new(Some(PlSmallStr::from_str("UTC")))
.expect("UTC timezone should always be valid")
.expect("UTC timezone should always be present")
}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn test_timezone_from_str() -> Result<(), ShellError> {
let tz = timezone_from_str("America/New_York", None)?;
assert_eq!(tz.to_string(), "America/New_York");
Ok(())
}
#[test]
fn test_timezone_utc() {
let tz = timezone_utc();
assert_eq!(tz.to_string(), "UTC");
}
}

26
crates/nu_plugin_polars/src/dataframe/command/datetime/replace_time_zone.rs
View File

@ -7,12 +7,15 @@ use crate::{
use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand};
use nu_protocol::{
Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type,
Value,
Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, Spanned,
SyntaxShape, Type, Value,
};
use chrono::DateTime;
use polars::prelude::*;
use polars_plan::plans::DynLiteralValue;
use super::timezone_from_str;
#[derive(Clone)]
pub struct ReplaceTimeZone;
@ -93,7 +96,8 @@ impl PluginCommand for ReplaceTimeZone {
"datetime".into(),
DataType::Datetime(
TimeUnit::Nanoseconds,
Some(PlSmallStr::from_static("America/New_York")),
TimeZone::opt_try_new(Some("America/New_York"))
.expect("timezone should be valid"),
),
),
])))),
@ -145,7 +149,8 @@ impl PluginCommand for ReplaceTimeZone {
"datetime".into(),
DataType::Datetime(
TimeUnit::Nanoseconds,
Some(PlSmallStr::from_static("America/New_York")),
TimeZone::opt_try_new(Some("America/New_York"))
.expect("timezone should be valid"),
),
),
])))),
@ -197,7 +202,8 @@ impl PluginCommand for ReplaceTimeZone {
"datetime".into(),
DataType::Datetime(
TimeUnit::Nanoseconds,
Some(PlSmallStr::from_static("America/New_York")),
TimeZone::opt_try_new(Some("America/New_York"))
.expect("timezone should be valid"),
),
),
])))),
@ -255,13 +261,17 @@ impl PluginCommand for ReplaceTimeZone {
match PolarsPluginObject::try_from_value(plugin, &value)? {
PolarsPluginObject::NuExpression(expr) => {
let time_zone: String = call.req(0)?;
let time_zone_spanned: Spanned<String> = call.req(0)?;
let time_zone =
timezone_from_str(&time_zone_spanned.item, Some(time_zone_spanned.span))?;
let expr: NuExpression = expr
.into_polars()
.dt()
.replace_time_zone(
Some(PlSmallStr::from_str(&time_zone)),
Expr::Literal(LiteralValue::String(PlSmallStr::from_string(ambiguous))),
Some(time_zone),
Expr::Literal(LiteralValue::Dyn(DynLiteralValue::Str(
PlSmallStr::from_string(ambiguous),
))),
nonexistent,
)
.into();

5
crates/nu_plugin_polars/src/dataframe/command/datetime/truncate.rs
View File

@ -15,6 +15,7 @@ use nu_protocol::{
use chrono::DateTime;
use polars::prelude::{DataType, Expr, Field, LiteralValue, PlSmallStr, Schema, TimeUnit};
use polars_plan::plans::DynLiteralValue;
#[derive(Clone)]
pub struct Truncate;
@ -191,9 +192,9 @@ fn command(
let res: NuExpression = expr
.into_polars()
.dt()
.truncate(Expr::Literal(LiteralValue::String(
.truncate(Expr::Literal(LiteralValue::Dyn(DynLiteralValue::Str(
PlSmallStr::from_string(every),
)))
))))
.into();
res.to_pipeline_data(plugin, engine, call.head)
}

6
crates/nu_plugin_polars/src/dataframe/command/list/contains.rs
View File

@ -147,7 +147,11 @@ fn command_expr(
});
}
};
let res: NuExpression = expr.into_polars().list().contains(single_expression).into();
let res: NuExpression = expr
.into_polars()
.list()
.contains(single_expression, true)
.into();
res.to_pipeline_data(plugin, engine, call.head)
}

6
crates/nu_plugin_polars/src/dataframe/values/mod.rs
View File

@ -417,6 +417,8 @@ pub trait CustomValueSupport: Cacheable {
mod test {
use polars::prelude::{DataType, TimeUnit, UnknownKind};
use crate::command::datetime::timezone_utc;
use super::*;
#[test]
@ -498,7 +500,7 @@ mod test {
let dtype = "object";
let schema = str_to_dtype(dtype, Span::unknown()).unwrap();
let expected = DataType::Object("unknown", None);
let expected = DataType::Object("unknown");
assert_eq!(schema, expected);
}
@ -526,7 +528,7 @@ mod test {
let dtype = "datetime<ms, UTC>";
let schema = str_to_dtype(dtype, Span::unknown()).unwrap();
let expected = DataType::Datetime(TimeUnit::Milliseconds, Some("UTC".into()));
let expected = DataType::Datetime(TimeUnit::Milliseconds, Some(timezone_utc()));
assert_eq!(schema, expected);
let dtype = "invalid";

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

@ -14,13 +14,15 @@ use polars::prelude::{
Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, IntoSeries, ListBooleanChunkedBuilder,
ListBuilderTrait, ListPrimitiveChunkedBuilder, ListStringChunkedBuilder, ListType, LogicalType,
NamedFrom, NewChunkedArray, ObjectType, PolarsError, Schema, SchemaExt, Series, StructChunked,
TemporalMethods, TimeUnit, UInt8Type, UInt16Type, UInt32Type, UInt64Type,
TemporalMethods, TimeUnit, TimeZone as PolarsTimeZone, UInt8Type, UInt16Type, UInt32Type,
UInt64Type,
};
use nu_protocol::{Record, ShellError, Span, Value};
use polars_arrow::Either;
use polars_arrow::array::Utf8ViewArray;
use crate::command::datetime::timezone_utc;
use crate::dataframe::values::NuSchema;
use super::{DataFrameValue, NuDataFrame};
@ -232,7 +234,7 @@ pub fn insert_value(
col_val.column_type = value_to_data_type(&value);
} else if let Some(current_data_type) = current_data_type {
if col_val.column_type.as_ref() != Some(&current_data_type) {
col_val.column_type = Some(DataType::Object("Value", None));
col_val.column_type = Some(DataType::Object("Value"));
}
}
col_val.values.push(value);
@ -248,7 +250,7 @@ fn value_to_data_type(value: &Value) -> Option<DataType> {
Value::Bool { .. } => Some(DataType::Boolean),
Value::Date { .. } => Some(DataType::Datetime(
TimeUnit::Nanoseconds,
Some(PlSmallStr::from_static("UTC")),
Some(timezone_utc()),
)),
Value::Duration { .. } => Some(DataType::Duration(TimeUnit::Nanoseconds)),
Value::Filesize { .. } => Some(DataType::Int64),
@ -266,7 +268,7 @@ fn value_to_data_type(value: &Value) -> Option<DataType> {
.map(value_to_data_type)
.nth(1)
.flatten()
.unwrap_or(DataType::Object("Value", None));
.unwrap_or(DataType::Object("Value"));
Some(DataType::List(Box::new(list_type)))
}
@ -278,7 +280,7 @@ fn typed_column_to_series(name: PlSmallStr, column: TypedColumn) -> Result<Serie
let column_type = &column
.column_type
.clone()
.unwrap_or(DataType::Object("Value", None));
.unwrap_or(DataType::Object("Value"));
match column_type {
DataType::Float32 => {
let series_values: Result<Vec<_>, _> = column
@ -433,7 +435,7 @@ fn typed_column_to_series(name: PlSmallStr, column: TypedColumn) -> Result<Serie
column.values.iter().map(|v| v.coerce_binary()).collect();
Ok(Series::new(name, series_values?))
}
DataType::Object(_, _) => value_to_series(name, &column.values),
DataType::Object(_) => value_to_series(name, &column.values),
DataType::Duration(time_unit) => {
let series_values: Result<Vec<_>, _> = column
.values
@ -452,11 +454,7 @@ fn typed_column_to_series(name: PlSmallStr, column: TypedColumn) -> Result<Serie
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,
&DataType::Object("unknown", None),
&column.values,
)
input_type_list_to_series(&name, &DataType::Object("unknown"), &column.values)
}
}
}
@ -1108,7 +1106,7 @@ fn series_to_values(
Ok(values)
}
DataType::Object(x, _) => {
DataType::Object(x) => {
let casted = series
.as_any()
.downcast_ref::<ChunkedArray<ObjectType<DataFrameValue>>>();
@ -1451,7 +1449,7 @@ fn nanos_to_timeunit(a: i64, time_unit: TimeUnit) -> Result<i64, ShellError> {
fn datetime_from_epoch_nanos(
nanos: i64,
timezone: &Option<PlSmallStr>,
timezone: &Option<PolarsTimeZone>,
span: Span,
) -> Result<DateTime<FixedOffset>, ShellError> {
let tz: Tz = if let Some(polars_tz) = timezone {

14
crates/nu_plugin_polars/src/dataframe/values/nu_dtype/mod.rs
View File

@ -2,10 +2,10 @@ pub mod custom_value;
use custom_value::NuDataTypeCustomValue;
use nu_protocol::{ShellError, Span, Value, record};
use polars::prelude::{DataType, Field, PlSmallStr, TimeUnit, UnknownKind};
use polars::prelude::{DataType, Field, TimeUnit, UnknownKind};
use uuid::Uuid;
use crate::{Cacheable, PolarsPlugin};
use crate::{Cacheable, PolarsPlugin, command::datetime::timezone_from_str};
use super::{CustomValueSupport, PolarsPluginObject, PolarsPluginType};
@ -166,7 +166,7 @@ pub fn str_to_dtype(dtype: &str, span: Span) -> Result<DataType, ShellError> {
"time" => Ok(DataType::Time),
"null" => Ok(DataType::Null),
"unknown" => Ok(DataType::Unknown(UnknownKind::Any)),
"object" => Ok(DataType::Object("unknown", None)),
"object" => Ok(DataType::Object("unknown")),
_ if dtype.starts_with("list") => {
let dtype = dtype
.trim_start_matches("list")
@ -206,12 +206,10 @@ pub fn str_to_dtype(dtype: &str, span: Span) -> Result<DataType, ShellError> {
let timezone = if "*" == next {
None
} else {
Some(next.to_string())
let zone_str = next.to_string();
Some(timezone_from_str(&zone_str, None)?)
};
Ok(DataType::Datetime(
time_unit,
timezone.map(PlSmallStr::from),
))
Ok(DataType::Datetime(time_unit, timezone))
}
_ if dtype.starts_with("duration") => {
let inner = dtype.trim_start_matches("duration<").trim_end_matches('>');

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

@ -252,7 +252,7 @@ pub fn expr_to_value(expr: &Expr, span: Span) -> Result<Value, ShellError> {
record! { "expr" => Value::string("wildcard", span) },
span,
)),
Expr::Explode(expr) => Ok(Value::record(
Expr::Explode { input: expr, .. } => Ok(Value::record(
record! { "expr" => expr_to_value(expr.as_ref(), span)? },
span,
)),