Added command polars len for performing count(*) like operations. (#13941)

# Description
This request exposes the prelude::polars::len expression. It is ended
for doing fast select count(*) like operations:

<img width="626" alt="Screenshot 2024-09-26 at 18 14 45"
src="https://github.com/user-attachments/assets/74285fc6-f99c-46e0-9226-9a7d41738d78">

# User-Facing Changes
- Introduction of the `polars len` command
This commit is contained in:
Jack Wright 2024-09-27 04:54:28 -07:00 committed by GitHub
parent d68c3ec89a
commit 5bef81a059
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 78 additions and 0 deletions

View file

@ -0,0 +1,76 @@
use crate::{
values::{Column, CustomValueSupport, NuDataFrame, NuExpression},
PolarsPlugin,
};
use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand};
use nu_protocol::{Category, Example, LabeledError, PipelineData, Signature, Span, Type, Value};
#[derive(Clone)]
pub struct ExprLen;
impl PluginCommand for ExprLen {
type Plugin = PolarsPlugin;
fn name(&self) -> &str {
"polars len"
}
fn description(&self) -> &str {
"Return the number of rows in the context. This is similar to COUNT(*) in SQL."
}
fn signature(&self) -> Signature {
Signature::build(self.name())
.input_output_type(Type::Any, Type::Custom("expression".into()))
.category(Category::Custom("dataframe".into()))
}
fn examples(&self) -> Vec<Example> {
vec![
Example {
description: "Count the number of rows in the the dataframe.",
example: "[[a b]; [1 2] [3 4]] | polars into-df | polars select (polars len) | polars collect",
result: Some(
NuDataFrame::try_from_columns(
vec![
Column::new("len".to_string(), vec![Value::test_int(2)]),
],
None,
)
.expect("simple df for test should not fail")
.into_value(Span::test_data()),
),
},
Example {
description: "Creates a last expression from a column",
example: "polars col a | polars last",
result: None,
},
]
}
fn run(
&self,
plugin: &Self::Plugin,
engine: &EngineInterface,
call: &EvaluatedCall,
_input: PipelineData,
) -> Result<PipelineData, LabeledError> {
let res: NuExpression = polars::prelude::len().into();
res.to_pipeline_data(plugin, engine, call.head)
.map_err(LabeledError::from)
}
}
#[cfg(test)]
mod test {
use super::*;
use crate::test::test_polars_plugin_command;
use nu_protocol::ShellError;
#[test]
fn test_examples() -> Result<(), ShellError> {
test_polars_plugin_command(&ExprLen)
}
}

View file

@ -19,6 +19,7 @@ mod flatten;
mod get; mod get;
mod join; mod join;
mod last; mod last;
mod len;
mod lit; mod lit;
mod pivot; mod pivot;
mod query_df; mod query_df;
@ -85,6 +86,7 @@ pub(crate) fn data_commands() -> Vec<Box<dyn PluginCommand<Plugin = PolarsPlugin
Box::new(UnpivotDF), Box::new(UnpivotDF),
Box::new(FirstDF), Box::new(FirstDF),
Box::new(LastDF), Box::new(LastDF),
Box::new(len::ExprLen),
Box::new(RenameDF), Box::new(RenameDF),
Box::new(SampleDF), Box::new(SampleDF),
Box::new(SliceDF), Box::new(SliceDF),