diff --git a/crates/nu_plugin_polars/src/dataframe/eager/to_df.rs b/crates/nu_plugin_polars/src/dataframe/eager/to_df.rs index 62f662e2a8..0747f16613 100644 --- a/crates/nu_plugin_polars/src/dataframe/eager/to_df.rs +++ b/crates/nu_plugin_polars/src/dataframe/eager/to_df.rs @@ -37,6 +37,11 @@ impl PluginCommand for ToDataFrame { r#"Polars Schema in format [{name: str}]. CSV, JSON, and JSONL files"#, Some('s'), ) + .switch( + "as-columns", + r#"When input shape is record of lists, treat each list as column values."#, + Some('c'), + ) .input_output_type(Type::Any, Type::Custom("dataframe".into())) .category(Category::Custom("dataframe".into())) } @@ -64,6 +69,27 @@ impl PluginCommand for ToDataFrame { .into_value(Span::test_data()), ), }, + Example { + description: "Takes a record of lists and creates a dataframe", + example: "{a: [1 3], b: [2 4]} | polars into-df --as-columns", + result: Some( + NuDataFrame::try_from_columns( + vec![ + Column::new( + "a".to_string(), + vec![Value::test_int(1), Value::test_int(3)], + ), + Column::new( + "b".to_string(), + vec![Value::test_int(2), Value::test_int(4)], + ), + ], + None, + ) + .expect("simple df for test should not fail") + .into_value(Span::test_data()), + ), + }, Example { description: "Takes a list of tables and creates a dataframe", example: "[[1 2 a] [3 4 b] [5 6 c]] | polars into-df", @@ -182,7 +208,39 @@ impl PluginCommand for ToDataFrame { .map(|schema| NuSchema::try_from(&schema)) .transpose()?; - let df = NuDataFrame::try_from_iter(plugin, input.into_iter(), maybe_schema.clone())?; + let maybe_as_columns = call.has_flag("as-columns")?; + + let df = if !maybe_as_columns { + NuDataFrame::try_from_iter(plugin, input.into_iter(), maybe_schema.clone())? + } else { + match &input { + PipelineData::Value(Value::Record { val, .. }, _) => { + let items: Result)>, &str> = val + .iter() + .map(|(k, v)| match v.to_owned().into_list() { + Ok(v) => Ok((k.to_owned(), v)), + _ => Err("error"), + }) + .collect(); + match items { + Ok(items) => { + let columns = items + .iter() + .map(|(k, v)| Column::new(k.to_owned(), v.to_owned())) + .collect::>(); + NuDataFrame::try_from_columns(columns, maybe_schema)? + } + Err(_) => NuDataFrame::try_from_iter( + plugin, + input.into_iter(), + maybe_schema.clone(), + )?, + } + } + _ => NuDataFrame::try_from_iter(plugin, input.into_iter(), maybe_schema.clone())?, + } + }; + df.to_pipeline_data(plugin, engine, call.head) .map_err(LabeledError::from) } diff --git a/crates/nu_plugin_polars/src/dataframe/eager/to_nu.rs b/crates/nu_plugin_polars/src/dataframe/eager/to_nu.rs index 72aff7a728..0647b8d4e3 100644 --- a/crates/nu_plugin_polars/src/dataframe/eager/to_nu.rs +++ b/crates/nu_plugin_polars/src/dataframe/eager/to_nu.rs @@ -23,7 +23,7 @@ impl PluginCommand for ToNu { } fn usage(&self) -> &str { - "Converts a dataframe or an expression into into nushell value for access and exploration." + "Converts a dataframe or an expression into nushell value for access and exploration." } fn signature(&self) -> Signature {