Dataframe Shape command (#3805)

* size command to get dataframe info

* change command name to shape

* apply lint to file

crates/nu-command/src/commands/dataframe/mod.rs

@@ -18,6 +18,7 @@ pub mod open;
 pub mod pivot;
 pub mod sample;
 pub mod select;
+pub mod shape;
 pub mod show;
 pub mod slice;
 pub mod sort;
@@ -50,6 +51,7 @@ pub use open::DataFrame as DataFrameOpen;
 pub use pivot::DataFrame as DataFramePivot;
 pub use sample::DataFrame as DataFrameSample;
 pub use select::DataFrame as DataFrameSelect;
+pub use shape::DataFrame as DataFrameShape;
 pub use show::DataFrame as DataFrameShow;
 pub use slice::DataFrame as DataFrameSlice;
 pub use sort::DataFrame as DataFrameSort;

crates/nu-command/src/commands/dataframe/shape.rs

@@ -0,0 +1,47 @@
+use crate::prelude::*;
+use nu_engine::WholeStreamCommand;
+use nu_errors::ShellError;
+use nu_protocol::{dataframe::NuDataFrame, Signature, TaggedDictBuilder};
+
+pub struct DataFrame;
+
+impl WholeStreamCommand for DataFrame {
+    fn name(&self) -> &str {
+        "dataframe shape"
+    }
+
+    fn usage(&self) -> &str {
+        "[DataFrame] Shows column and row size for a dataframe"
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::build("dataframe shape")
+    }
+
+    fn run(&self, args: CommandArgs) -> Result<OutputStream, ShellError> {
+        command(args)
+    }
+
+    fn examples(&self) -> Vec<Example> {
+        vec![Example {
+            description: "Shows row and column shape",
+            example: "[[a b]; [1 2] [3 4]] | dataframe to-df | dataframe shape",
+            result: None,
+        }]
+    }
+}
+
+fn command(mut args: CommandArgs) -> Result<OutputStream, ShellError> {
+    let tag = args.call_info.name_tag.clone();
+
+    let df = NuDataFrame::try_from_stream(&mut args.input, &tag.span)?;
+
+    let rows = df.as_ref().height();
+    let cols = df.as_ref().width();
+
+    let mut data = TaggedDictBuilder::new(&tag);
+    data.insert_value("rows", format!("{}", rows));
+    data.insert_value("columns", format!("{}", cols));
+
+    Ok(OutputStream::one(data.into_value()))
+}

crates/nu-command/src/commands/mod.rs

@@ -32,9 +32,10 @@ pub use dataframe::{
     DataFrameIsIn, DataFrameIsNotNull, DataFrameIsNull, DataFrameIsUnique, DataFrameJoin,
     DataFrameLast, DataFrameList, DataFrameMelt, DataFrameNNull, DataFrameNUnique, DataFrameNot,
     DataFrameOpen, DataFramePivot, DataFrameSample, DataFrameSelect, DataFrameSeriesRename,
-    DataFrameSet, DataFrameSetWithIdx, DataFrameShift, DataFrameShow, DataFrameSlice,
+    DataFrameSet, DataFrameSetWithIdx, DataFrameShape, DataFrameShift, DataFrameShow,
-    DataFrameSort, DataFrameTake, DataFrameToCsv, DataFrameToDF, DataFrameToParquet,
+    DataFrameSlice, DataFrameSort, DataFrameTake, DataFrameToCsv, DataFrameToDF,
-    DataFrameToSeries, DataFrameUnique, DataFrameValueCounts, DataFrameWhere, DataFrameWithColumn,
+    DataFrameToParquet, DataFrameToSeries, DataFrameUnique, DataFrameValueCounts, DataFrameWhere,
+    DataFrameWithColumn,
 };
 pub use env::*;
 pub use filesystem::*;

crates/nu-command/src/default_context.rs

@@ -317,6 +317,7 @@ pub fn create_default_context(interactive: bool) -> Result<EvaluationContext, Bo
             whole_stream_command(DataFrameNot),
             whole_stream_command(DataFrameTake),
             whole_stream_command(DataFrameSetWithIdx),
+            whole_stream_command(DataFrameShape),
         ]);
 
         #[cfg(feature = "clipboard-cli")]