extern/nushell
1
0
Fork 1
mirror of https://github.com/nushell/nushell.git synced 2025-06-30 14:40:06 +02:00
Code Issues Packages Projects Releases Wiki Activity

Polars Struct support without unsafe blocks (#11229)

Browse Source 
Second attempt at polars Struct support. This version avoid using unsafe
checks by cloning the StructArray and utilizing the into_static to
convert to a StructOwned.

---------

Co-authored-by: Jack Wright <jack.wright@disqo.com>
This commit is contained in:
Jack Wright
2023-12-15 02:21:30 -08:00
committed by GitHub
parent 6ead98effb
commit 44dc890124
2 changed files with 474 additions and 195 deletions
Download Patch File Download Diff File Expand all files Collapse all files

505
crates/nu-cmd-dataframe/src/dataframe/values/nu_dataframe/conversion.rs
View File

@ -1,11 +1,12 @@
use super::{DataFrameValue, NuDataFrame};
use std::ops::{Deref, DerefMut};
use chrono::{DateTime, FixedOffset, NaiveDateTime};
use chrono::{DateTime, Duration, FixedOffset, NaiveTime, TimeZone, Utc};
use chrono_tz::Tz;
use indexmap::map::{Entry, IndexMap};
use nu_protocol::{Record, ShellError, Span, Value};
use polars::chunked_array::builder::AnonymousOwnedListBuilder;
use polars::chunked_array::object::builder::ObjectChunkedBuilder;
use polars::chunked_array::ChunkedArray;
use polars::datatypes::AnyValue;
use polars::export::arrow::Either;
use polars::prelude::{
DataFrame, DataType, DatetimeChunked, Float64Type, Int64Type, IntoSeries,
@ -13,11 +14,14 @@ use polars::prelude::{
ListUtf8ChunkedBuilder, NamedFrom, NewChunkedArray, ObjectType, Series, TemporalMethods,
TimeUnit,
};
use std::ops::{Deref, DerefMut};
const SECS_PER_DAY: i64 = 86_400;
use nu_protocol::{Record, ShellError, Span, Value};
// The values capacity is for the size of an internal vec.
use super::{DataFrameValue, NuDataFrame};
const NANOS_PER_DAY: i64 = 86_400_000_000_000;
// The values capacity is for the size of an vec.
// Since this is impossible to determine without traversing every value
// I just picked one. Since this is for converting back and forth
// between nushell tables the values shouldn't be too extremely large for
@ -199,7 +203,7 @@ fn value_to_input_type(value: &Value) -> InputType {
Value::Filesize { .. } => InputType::Filesize,
Value::List { vals, .. } => {
// We need to determined the type inside of the list.
// Since Value::List does not have any kind of internal
// Since Value::List does not have any kind of
// type information, we need to look inside the list.
// This will cause errors if lists have inconsistent types.
// Basically, if a list column needs to be converted to dataframe,
@ -775,28 +779,21 @@ fn series_to_values(
}),
Some(ca) => {
let it = ca.into_iter();
let values: Vec<Value> =
if let (Some(size), Some(from_row)) = (maybe_size, maybe_from_row) {
Either::Left(it.skip(from_row).take(size))
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(|ca| {
let sublist: Vec<Value> = if let Some(ref s) = ca {
series_to_values(s, None, None, Span::unknown())?
} else {
Either::Right(it)
}
.map(|ca| {
let sublist = ca
.map(|ref s| {
match series_to_values(s, None, None, Span::unknown()) {
Ok(v) => v,
Err(e) => {
eprintln!("Error list values: {e}");
vec![]
}
}
})
.unwrap_or(vec![]);
Value::list(sublist, span)
})
.collect::<Vec<Value>>();
Ok(values)
// empty item
vec![]
};
Ok(Value::list(sublist, span))
})
.collect::<Result<Vec<Value>, ShellError>>()
}
}
}
@ -817,51 +814,16 @@ fn series_to_values(
}
.map(|v| match v {
Some(a) => {
// elapsed time in day since 1970-01-01
let seconds = a as i64 * SECS_PER_DAY;
let naive_datetime = match NaiveDateTime::from_timestamp_opt(seconds, 0) {
Some(val) => val,
None => {
return Value::error(
ShellError::UnsupportedInput {
msg: "The given local datetime representation is invalid."
.to_string(),
input: format!("timestamp is {a:?}"),
msg_span: span,
input_span: Span::unknown(),
},
span,
)
}
};
// Zero length offset
let offset = match FixedOffset::east_opt(0) {
Some(val) => val,
None => {
return Value::error(
ShellError::UnsupportedInput {
msg: "The given local datetime representation is invalid."
.to_string(),
input: format!("timestamp is {a:?}"),
msg_span: span,
input_span: Span::unknown(),
},
span,
)
}
};
let datetime =
DateTime::<FixedOffset>::from_naive_utc_and_offset(naive_datetime, offset);
Value::date(datetime, span)
let nanos = nanos_per_day(a);
let datetime = datetime_from_epoch_nanos(nanos, &None, span)?;
Ok(Value::date(datetime, span))
}
None => Value::nothing(span),
None => Ok(Value::nothing(span)),
})
.collect::<Vec<Value>>();
.collect::<Result<Vec<Value>, ShellError>>()?;
Ok(values)
}
DataType::Datetime(time_unit, _) => {
DataType::Datetime(time_unit, tz) => {
let casted = series.datetime().map_err(|e| ShellError::GenericError {
error: "Error casting column to datetime".into(),
msg: "".into(),
@ -878,55 +840,46 @@ fn series_to_values(
}
.map(|v| match v {
Some(a) => {
let unit_divisor = match time_unit {
TimeUnit::Nanoseconds => 1_000_000_000,
TimeUnit::Microseconds => 1_000_000,
TimeUnit::Milliseconds => 1_000,
};
// elapsed time in nano/micro/milliseconds since 1970-01-01
let seconds = a / unit_divisor;
let naive_datetime = match NaiveDateTime::from_timestamp_opt(seconds, 0) {
Some(val) => val,
None => {
return Value::error(
ShellError::UnsupportedInput {
msg: "The given local datetime representation is invalid."
.to_string(),
input: format!("timestamp is {a:?}"),
msg_span: span,
input_span: Span::unknown(),
},
span,
)
}
};
// Zero length offset
let offset = match FixedOffset::east_opt(0) {
Some(val) => val,
None => {
return Value::error(
ShellError::UnsupportedInput {
msg: "The given local datetime representation is invalid."
.to_string(),
input: format!("timestamp is {a:?}"),
msg_span: span,
input_span: Span::unknown(),
},
span,
)
}
};
let datetime =
DateTime::<FixedOffset>::from_naive_utc_and_offset(naive_datetime, offset);
Value::date(datetime, span)
let nanos = nanos_from_timeunit(a, *time_unit);
let datetime = datetime_from_epoch_nanos(nanos, tz, span)?;
Ok(Value::date(datetime, span))
}
None => Value::nothing(span),
None => Ok(Value::nothing(span)),
})
.collect::<Vec<Value>>();
.collect::<Result<Vec<Value>, ShellError>>()?;
Ok(values)
}
DataType::Struct(polar_fields) => {
let casted = series.struct_().map_err(|e| ShellError::GenericError {
error: "Error casting column to struct".into(),
msg: "".to_string(),
span: None,
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 vals: Result<Vec<Value>, ShellError> = any_values
.iter()
.map(|v| any_value_to_value(v, span))
.collect();
let cols: Vec<String> = polar_fields
.iter()
.map(|field| field.name.to_string())
.collect();
let record = Record { cols, vals: vals? };
Ok(Value::record(record, span))
})
.collect();
values
}
DataType::Time => {
let casted =
series
@ -963,10 +916,154 @@ fn series_to_values(
}
}
fn any_value_to_value(any_value: &AnyValue, span: Span) -> Result<Value, ShellError> {
match any_value {
AnyValue::Null => Ok(Value::nothing(span)),
AnyValue::Boolean(b) => Ok(Value::bool(*b, span)),
AnyValue::Utf8(s) => Ok(Value::string(s.to_string(), span)),
AnyValue::UInt8(i) => Ok(Value::int(*i as i64, span)),
AnyValue::UInt16(i) => Ok(Value::int(*i as i64, span)),
AnyValue::UInt32(i) => Ok(Value::int(*i as i64, span)),
AnyValue::UInt64(i) => Ok(Value::int(*i as i64, span)),
AnyValue::Int8(i) => Ok(Value::int(*i as i64, span)),
AnyValue::Int16(i) => Ok(Value::int(*i as i64, span)),
AnyValue::Int32(i) => Ok(Value::int(*i as i64, span)),
AnyValue::Int64(i) => Ok(Value::int(*i, span)),
AnyValue::Float32(f) => Ok(Value::float(*f as f64, span)),
AnyValue::Float64(f) => Ok(Value::float(*f, span)),
AnyValue::Date(d) => {
let nanos = nanos_per_day(*d);
datetime_from_epoch_nanos(nanos, &None, span)
.map(|datetime| Value::date(datetime, span))
}
AnyValue::Datetime(a, time_unit, tz) => {
let nanos = nanos_from_timeunit(*a, *time_unit);
datetime_from_epoch_nanos(nanos, tz, span).map(|datetime| Value::date(datetime, span))
}
AnyValue::Duration(a, time_unit) => {
let nanos = match time_unit {
TimeUnit::Nanoseconds => *a,
TimeUnit::Microseconds => *a * 1_000,
TimeUnit::Milliseconds => *a * 1_000_000,
};
Ok(Value::duration(nanos, span))
}
// AnyValue::Time represents the current time since midnight.
// Unfortunately, there is no timezone related information.
// Given this, calculate the current date from UTC and add the time.
AnyValue::Time(nanos) => time_from_midnight(*nanos, span),
AnyValue::List(series) => {
series_to_values(series, None, None, span).map(|values| Value::list(values, span))
}
AnyValue::Struct(_idx, _struct_array, _s_fields) => {
// This should convert to a StructOwned object.
let static_value =
any_value
.clone()
.into_static()
.map_err(|e| ShellError::GenericError {
error: "Cannot convert polars struct to static value".into(),
msg: e.to_string(),
span: Some(span),
help: None,
inner: Vec::new(),
})?;
any_value_to_value(&static_value, span)
}
AnyValue::StructOwned(struct_tuple) => {
let values: Result<Vec<Value>, ShellError> = struct_tuple
.0
.iter()
.map(|s| any_value_to_value(s, span))
.collect();
let fields = struct_tuple
.1
.iter()
.map(|f| f.name().to_string())
.collect();
Ok(Value::Record {
val: Record {
cols: fields,
vals: values?,
},
internal_span: span,
})
}
AnyValue::Utf8Owned(s) => Ok(Value::string(s.to_string(), span)),
AnyValue::Binary(bytes) => Ok(Value::binary(*bytes, span)),
AnyValue::BinaryOwned(bytes) => Ok(Value::binary(bytes.to_owned(), span)),
e => Err(ShellError::GenericError {
error: "Error creating Value".into(),
msg: "".to_string(),
span: None,
help: Some(format!("Value not supported in nushell: {e}")),
inner: Vec::new(),
}),
}
}
fn nanos_per_day(days: i32) -> i64 {
days as i64 * NANOS_PER_DAY
}
fn nanos_from_timeunit(a: i64, time_unit: TimeUnit) -> i64 {
a * match time_unit {
TimeUnit::Microseconds => 1_000, // Convert microseconds to nanoseconds
TimeUnit::Milliseconds => 1_000_000, // Convert milliseconds to nanoseconds
TimeUnit::Nanoseconds => 1, // Already in nanoseconds
}
}
fn datetime_from_epoch_nanos(
nanos: i64,
timezone: &Option<String>,
span: Span,
) -> Result<DateTime<FixedOffset>, ShellError> {
let tz: Tz = if let Some(polars_tz) = timezone {
polars_tz
.parse::<Tz>()
.map_err(|_| ShellError::GenericError {
error: format!("Could not parse polars timezone: {polars_tz}"),
msg: "".to_string(),
span: Some(span),
help: None,
inner: vec![],
})?
} else {
Tz::UTC
};
Ok(tz.timestamp_nanos(nanos).fixed_offset())
}
fn time_from_midnight(nanos: i64, span: Span) -> Result<Value, ShellError> {
let today = Utc::now().date_naive();
NaiveTime::from_hms_opt(0, 0, 0) // midnight
.map(|time| time + Duration::nanoseconds(nanos)) // current time
.map(|time| today.and_time(time)) // current date and time
.and_then(|datetime| {
FixedOffset::east_opt(0) // utc
.map(|offset| {
DateTime::<FixedOffset>::from_naive_utc_and_offset(datetime, offset)
})
})
.map(|datetime| Value::date(datetime, span)) // current date and time
.ok_or(ShellError::CantConvert {
to_type: "datetime".to_string(),
from_type: "polars time".to_string(),
span,
help: Some("Could not convert polars time of {nanos} to datetime".to_string()),
})
}
#[cfg(test)]
mod tests {
use super::*;
use indexmap::indexmap;
use polars::export::arrow::array::{BooleanArray, PrimitiveArray};
use polars::prelude::Field;
use polars_io::prelude::StructArray;
use super::*;
#[test]
fn test_parsed_column_string_list() -> Result<(), Box<dyn std::error::Error>> {
@ -1001,4 +1098,184 @@ mod tests {
Ok(())
}
#[test]
fn test_any_value_to_value() -> Result<(), Box<dyn std::error::Error>> {
let span = Span::test_data();
assert_eq!(
any_value_to_value(&AnyValue::Null, span)?,
Value::nothing(span)
);
let test_bool = true;
assert_eq!(
any_value_to_value(&AnyValue::Boolean(test_bool), span)?,
Value::bool(test_bool, span)
);
let test_str = "foo";
assert_eq!(
any_value_to_value(&AnyValue::Utf8(test_str), span)?,
Value::string(test_str.to_string(), span)
);
assert_eq!(
any_value_to_value(&AnyValue::Utf8Owned(test_str.into()), span)?,
Value::string(test_str.to_owned(), span)
);
let tests_uint8 = 4;
assert_eq!(
any_value_to_value(&AnyValue::UInt8(tests_uint8), span)?,
Value::int(tests_uint8 as i64, span)
);
let tests_uint16 = 233;
assert_eq!(
any_value_to_value(&AnyValue::UInt16(tests_uint16), span)?,
Value::int(tests_uint16 as i64, span)
);
let tests_uint32 = 897688233;
assert_eq!(
any_value_to_value(&AnyValue::UInt32(tests_uint32), span)?,
Value::int(tests_uint32 as i64, span)
);
let tests_uint64 = 903225135897388233;
assert_eq!(
any_value_to_value(&AnyValue::UInt64(tests_uint64), span)?,
Value::int(tests_uint64 as i64, span)
);
let tests_float32 = 903225135897388233.3223353;
assert_eq!(
any_value_to_value(&AnyValue::Float32(tests_float32), span)?,
Value::float(tests_float32 as f64, span)
);
let tests_float64 = 9064251358973882322333.64233533232;
assert_eq!(
any_value_to_value(&AnyValue::Float64(tests_float64), span)?,
Value::float(tests_float64, span)
);
let test_days = 10_957;
let comparison_date = Utc
.with_ymd_and_hms(2000, 1, 1, 0, 0, 0)
.unwrap()
.fixed_offset();
assert_eq!(
any_value_to_value(&AnyValue::Date(test_days), span)?,
Value::date(comparison_date, span)
);
let test_millis = 946_684_800_000;
assert_eq!(
any_value_to_value(
&AnyValue::Datetime(test_millis, TimeUnit::Milliseconds, &None),
span
)?,
Value::date(comparison_date, span)
);
let test_duration_millis = 99_999;
let test_duration_micros = 99_999_000;
let test_duration_nanos = 99_999_000_000;
assert_eq!(
any_value_to_value(
&AnyValue::Duration(test_duration_nanos, TimeUnit::Nanoseconds),
span
)?,
Value::duration(test_duration_nanos, span)
);
assert_eq!(
any_value_to_value(
&AnyValue::Duration(test_duration_micros, TimeUnit::Microseconds),
span
)?,
Value::duration(test_duration_nanos, span)
);
assert_eq!(
any_value_to_value(
&AnyValue::Duration(test_duration_millis, TimeUnit::Milliseconds),
span
)?,
Value::duration(test_duration_nanos, span)
);
let test_binary = b"sdf2332f32q3f3afwaf3232f32";
assert_eq!(
any_value_to_value(&AnyValue::Binary(test_binary), span)?,
Value::binary(test_binary.to_vec(), span)
);
assert_eq!(
any_value_to_value(&AnyValue::BinaryOwned(test_binary.to_vec()), span)?,
Value::binary(test_binary.to_vec(), span)
);
let test_time_nanos = 54_000_000_000_000;
let test_time = DateTime::<FixedOffset>::from_naive_utc_and_offset(
Utc::now()
.date_naive()
.and_time(NaiveTime::from_hms_opt(15, 00, 00).unwrap()),
FixedOffset::east_opt(0).unwrap(),
);
assert_eq!(
any_value_to_value(&AnyValue::Time(test_time_nanos), span)?,
Value::date(test_time, span)
);
let test_list_series = Series::new("int series", &[1, 2, 3]);
let comparison_list_series = Value::list(
vec![
Value::int(1, span),
Value::int(2, span),
Value::int(3, span),
],
span,
);
assert_eq!(
any_value_to_value(&AnyValue::List(test_list_series), span)?,
comparison_list_series
);
let field_value_0 = AnyValue::Int32(1);
let field_value_1 = AnyValue::Boolean(true);
let values = vec![field_value_0, field_value_1];
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),
];
let test_owned_struct = AnyValue::StructOwned(Box::new((values, fields.clone())));
let comparison_owned_record = Value::record(
Record {
cols: vec![field_name_0.to_owned(), field_name_1.to_owned()],
vals: vec![Value::int(1, span), Value::bool(true, span)],
},
span,
);
assert_eq!(
any_value_to_value(&test_owned_struct, span)?,
comparison_owned_record.clone()
);
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(),
vec![Box::new(test_int_arr), Box::new(test_bool_arr)],
None,
);
assert_eq!(
any_value_to_value(
&AnyValue::Struct(0, &test_struct_arr, fields.as_slice()),
span
)?,
comparison_owned_record
);
Ok(())
}
}