From c6016d76591775323d7894893b1b8306758b73e5 Mon Sep 17 00:00:00 2001 From: Jack Wright <56345+ayax79@users.noreply.github.com> Date: Tue, 24 Oct 2023 19:25:21 -0700 Subject: [PATCH] Dataframe support for small int types (#10828) Turned features to allow signed and unsigned 8 and 16 bit types. --------- Co-authored-by: Jack Wright --- crates/nu-cmd-dataframe/Cargo.toml | 4 + .../src/dataframe/eager/dummies.rs | 73 +++++-------------- 2 files changed, 24 insertions(+), 53 deletions(-) diff --git a/crates/nu-cmd-dataframe/Cargo.toml b/crates/nu-cmd-dataframe/Cargo.toml index 60c737c164..65108008ad 100644 --- a/crates/nu-cmd-dataframe/Cargo.toml +++ b/crates/nu-cmd-dataframe/Cargo.toml @@ -38,6 +38,10 @@ features = [ "dtype-categorical", "dtype-datetime", "dtype-struct", + "dtype-i8", + "dtype-i16", + "dtype-u8", + "dtype-u16", "dynamic_group_by", "ipc", "is_in", diff --git a/crates/nu-cmd-dataframe/src/dataframe/eager/dummies.rs b/crates/nu-cmd-dataframe/src/dataframe/eager/dummies.rs index 72dfdbace7..7e2ae698a7 100644 --- a/crates/nu-cmd-dataframe/src/dataframe/eager/dummies.rs +++ b/crates/nu-cmd-dataframe/src/dataframe/eager/dummies.rs @@ -1,10 +1,10 @@ -use super::super::values::{Column, NuDataFrame}; +use super::super::values::NuDataFrame; use nu_protocol::{ ast::Call, engine::{Command, EngineState, Stack}, - Category, Example, PipelineData, ShellError, Signature, Span, Type, Value, + Category, Example, PipelineData, ShellError, Signature, Span, Type, }; -use polars::prelude::DataFrameOps; +use polars::{prelude::*, series::Series}; #[derive(Clone)] pub struct Dummies; @@ -34,24 +34,15 @@ impl Command for Dummies { description: "Create new dataframe with dummy variables from a dataframe", example: "[[a b]; [1 2] [3 4]] | dfr into-df | dfr dummies", result: Some( - NuDataFrame::try_from_columns(vec![ - Column::new( - "a_1".to_string(), - vec![Value::test_int(1), Value::test_int(0)], - ), - Column::new( - "a_3".to_string(), - vec![Value::test_int(0), Value::test_int(1)], - ), - Column::new( - "b_2".to_string(), - vec![Value::test_int(1), Value::test_int(0)], - ), - Column::new( - "b_4".to_string(), - vec![Value::test_int(0), Value::test_int(1)], - ), - ]) + NuDataFrame::try_from_series( + 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]), + ], + Span::test_data(), + ) .expect("simple df for test should not fail") .into_value(Span::test_data()), ), @@ -60,38 +51,14 @@ impl Command for Dummies { description: "Create new dataframe with dummy variables from a series", example: "[1 2 2 3 3] | dfr into-df | dfr dummies", result: Some( - NuDataFrame::try_from_columns(vec![ - Column::new( - "0_1".to_string(), - vec![ - Value::test_int(1), - Value::test_int(0), - Value::test_int(0), - Value::test_int(0), - Value::test_int(0), - ], - ), - Column::new( - "0_2".to_string(), - vec![ - Value::test_int(0), - Value::test_int(1), - Value::test_int(1), - Value::test_int(0), - Value::test_int(0), - ], - ), - Column::new( - "0_3".to_string(), - vec![ - Value::test_int(0), - Value::test_int(0), - Value::test_int(0), - Value::test_int(1), - Value::test_int(1), - ], - ), - ]) + NuDataFrame::try_from_series( + 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]), + ], + Span::test_data(), + ) .expect("simple df for test should not fail") .into_value(Span::test_data()), ),