add list-contains

crates/nu_plugin_polars/src/dataframe/command/list/contains.rs

@@ -0,0 +1,161 @@
+use crate::{
+    values::{
+        cant_convert_err, CustomValueSupport, NuExpression, PolarsPluginObject, PolarsPluginType,
+    },
+    PolarsPlugin,
+};
+
+use super::super::super::values::{Column, NuDataFrame};
+
+use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand};
+use nu_protocol::{
+    Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, SyntaxShape, Type,
+    Value,
+};
+
+#[derive(Clone)]
+pub struct ListContains;
+
+impl PluginCommand for ListContains {
+    type Plugin = PolarsPlugin;
+
+    fn name(&self) -> &str {
+        "polars list-contains"
+    }
+
+    fn description(&self) -> &str {
+        "Checks if an element is contained in a list."
+    }
+
+    fn signature(&self) -> Signature {
+        Signature::build(self.name())
+            .required(
+                "element",
+                SyntaxShape::Any,
+                "Element to search for in the list",
+            )
+            .input_output_types(vec![(
+                Type::Custom("expression".into()),
+                Type::Custom("expression".into()),
+            )])
+            .category(Category::Custom("dataframe".into()))
+    }
+
+    fn examples(&self) -> Vec<Example> {
+        vec![
+            Example {
+                description: "Returns boolean indicating if a literal element was found in a list column",
+                example: "let df = [[a]; [[a,b,c]] [[b,c,d]] [[c,d,f]]] | polars into-df -s {a: list<str>};
+                let df2 = $df | polars with-column [(polars col a | polars list-contains (polars lit a) | polars as b)] | polars collect;
+                $df2.b",
+                result: Some(
+                    NuDataFrame::try_from_columns(
+                        vec![Column::new(
+                            "b".to_string(),
+                            vec![
+                                Value::test_bool(true),
+                                Value::test_bool(false),
+                                Value::test_bool(false),
+                            ],
+                        )],
+                        None,
+                    )
+                    .expect("simple df for test should not fail")
+                    .into_value(Span::test_data()),
+                ),
+            },
+            Example {
+                description: "Returns boolean indicating if an element from another column was found in a list column",
+                example: "let df = [[a, b]; [[a,b,c], a] [[b,c,d], f] [[c,d,f], f]] | polars into-df -s {a: list<str>, b: str};
+                let df2 = $df | polars with-column [(polars col a | polars list-contains b | polars as c)] | polars collect;
+                $df2.c",
+                result: Some(
+                    NuDataFrame::try_from_columns(
+                        vec![Column::new(
+                            "b".to_string(),
+                            vec![
+                                Value::test_bool(true),
+                                Value::test_bool(false),
+                                Value::test_bool(true),
+                            ],
+                        )],
+                        None,
+                    )
+                    .expect("simple df for test should not fail")
+                    .into_value(Span::test_data()),
+                ),
+            },
+            Example {
+                description: "Returns boolean indicating if an element from another expression was found in a list column",
+                example: "let df = [[a, b]; [[1,2,3], 4] [[2,4,1], 2] [[2,1,6], 3]] | polars into-df -s {a: list<i64>, b: i64};
+                let df2 = $df | polars with-column [(polars col a | polars list-contains ((polars col b) * 2) | polars as c)] | polars collect;
+                $df2.c",
+                result: Some(
+                    NuDataFrame::try_from_columns(
+                        vec![Column::new(
+                            "b".to_string(),
+                            vec![
+                                Value::test_bool(false),
+                                Value::test_bool(true),
+                                Value::test_bool(true),
+                            ],
+                        )],
+                        None,
+                    )
+                    .expect("simple df for test should not fail")
+                    .into_value(Span::test_data()),
+                ),
+            }
+        ]
+    }
+
+    fn run(
+        &self,
+        plugin: &Self::Plugin,
+        engine: &EngineInterface,
+        call: &EvaluatedCall,
+        input: PipelineData,
+    ) -> Result<PipelineData, LabeledError> {
+        let value = input.into_value(call.head)?;
+        match PolarsPluginObject::try_from_value(plugin, &value)? {
+            PolarsPluginObject::NuExpression(expr) => command_expr(plugin, engine, call, expr),
+            _ => Err(cant_convert_err(&value, &[PolarsPluginType::NuExpression])),
+        }
+        .map_err(LabeledError::from)
+    }
+}
+
+fn command_expr(
+    plugin: &PolarsPlugin,
+    engine: &EngineInterface,
+    call: &EvaluatedCall,
+    expr: NuExpression,
+) -> Result<PipelineData, ShellError> {
+    let element = call.req(0)?;
+    let expressions = NuExpression::extract_exprs(plugin, element)?;
+    let single_expression = match expressions.as_slice() {
+        [single] => single.clone(),
+        _ => {
+            return Err(ShellError::GenericError {
+                error: "Expected a single polars expression".into(),
+                msg: "Requires a single polars expressions or column name as argument".into(),
+                span: Some(call.head),
+                help: None,
+                inner: vec![],
+            })
+        }
+    };
+    let res: NuExpression = expr.into_polars().list().contains(single_expression).into();
+    res.to_pipeline_data(plugin, engine, call.head)
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::test::test_polars_plugin_command;
+
+    #[test]
+    fn test_examples() -> Result<(), ShellError> {
+        test_polars_plugin_command(&ListContains)
+    }
+}

crates/nu_plugin_polars/src/dataframe/command/list/mod.rs

@@ -0,0 +1,10 @@
+mod contains;
+
+use crate::PolarsPlugin;
+use nu_plugin::PluginCommand;
+
+pub use contains::ListContains;
+
+pub(crate) fn list_commands() -> Vec<Box<dyn PluginCommand<Plugin = PolarsPlugin>>> {
+    vec![Box::new(ListContains)]
+}

crates/nu_plugin_polars/src/dataframe/command/mod.rs

@@ -5,5 +5,6 @@ pub mod data;
 pub mod datetime;
 pub mod index;
 pub mod integer;
+pub mod list;
 pub mod string;
 pub mod stub;

crates/nu_plugin_polars/src/lib.rs

@@ -8,7 +8,7 @@ pub use cache::{Cache, Cacheable};
 use command::{
     aggregation::aggregation_commands, boolean::boolean_commands, core::core_commands,
     data::data_commands, datetime::datetime_commands, index::index_commands,
-    integer::integer_commands, string::string_commands, stub::PolarsCmd,
+    integer::integer_commands, list::list_commands, string::string_commands, stub::PolarsCmd,
 };
 use log::debug;
 use nu_plugin::{EngineInterface, Plugin, PluginCommand};
@@ -93,6 +93,7 @@ impl Plugin for PolarsPlugin {
         commands.append(&mut index_commands());
         commands.append(&mut integer_commands());
         commands.append(&mut string_commands());
+        commands.append(&mut list_commands());
 
         commands.append(&mut cache_commands());
         commands