fix: check input type for arr/list.contains

pola-rs · Jan 25, 2024 · 17e973a · 17e973a
1 parent da867b7
commit 17e973a
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/crates/polars-plan/src/dsl/function_expr/array.rs b/crates/polars-plan/src/dsl/function_expr/array.rs
@@ -183,6 +183,9 @@ pub(super) fn join(s: &[Series], ignore_nulls: bool) -> PolarsResult<Series> {
 pub(super) fn contains(s: &[Series]) -> PolarsResult<Series> {
     let array = &s[0];
     let item = &s[1];
+    polars_ensure!(matches!(array.dtype(), DataType::Array(_, _)),
+        SchemaMismatch: "invalid series dtype: expected `Array`, got `{}`", array.dtype(),
+    );
     Ok(is_in(item, array)?.with_name(array.name()).into_series())
 }
 

diff --git a/crates/polars-plan/src/dsl/function_expr/list.rs b/crates/polars-plan/src/dsl/function_expr/list.rs
@@ -219,7 +219,9 @@ impl From<ListFunction> for SpecialEq<Arc<dyn SeriesUdf>> {
 pub(super) fn contains(args: &mut [Series]) -> PolarsResult<Option<Series>> {
     let list = &args[0];
     let item = &args[1];
-
+    polars_ensure!(matches!(list.dtype(), DataType::List(_)),
+        SchemaMismatch: "invalid series dtype: expected `List`, got `{}`", list.dtype(),
+    );
     polars_ops::prelude::is_in(item, list).map(|mut ca| {
         ca.rename(list.name());
         Some(ca.into_series())