feat(rust, python): add arr.take expression (#6116)

ritchie46 · web-flow · commit 05447ab387c0 · 2023-01-08T12:49:15.000+01:00
diff --git a/polars/Cargo.toml b/polars/Cargo.toml
@@ -124,8 +124,9 @@ list_eval = ["polars-lazy/list_eval"]
 cumulative_eval = ["polars-lazy/cumulative_eval"]
 chunked_ids = ["polars-core/chunked_ids", "polars-lazy/chunked_ids", "polars-core/chunked_ids"]
 to_dummies = ["polars-ops/to_dummies"]
-bigidx = ["polars-core/bigidx", "polars-lazy/bigidx"]
+bigidx = ["polars-core/bigidx", "polars-lazy/bigidx", "polars-ops/big_idx"]
 list_to_struct = ["polars-ops/list_to_struct", "polars-lazy/list_to_struct"]
+list_take = ["polars-ops/list_take", "polars-lazy/list_take"]
 describe = ["polars-core/describe"]
 timezones = ["polars-core/timezones", "polars-lazy/timezones"]
 string_justify = ["polars-lazy/string_justify", "polars-ops/string_justify"]
diff --git a/polars/polars-arrow/src/index.rs b/polars/polars-arrow/src/index.rs
@@ -2,22 +2,27 @@
 use arrow::array::UInt32Array;
 #[cfg(feature = "bigidx")]
 use arrow::array::UInt64Array;
+use num::{NumCast, Signed, Zero};
 
 pub trait IndexToUsize {
     /// Translate the negative index to an offset.
-    fn negative_to_usize(self, index: usize) -> Option<usize>;
+    fn negative_to_usize(self, len: usize) -> Option<usize>;
 }
 
-impl IndexToUsize for i64 {
-    fn negative_to_usize(self, index: usize) -> Option<usize> {
-        if self >= 0 && (self as usize) < index {
-            Some(self as usize)
+impl<I> IndexToUsize for I
+where
+    I: PartialOrd + PartialEq + NumCast + Signed + Zero,
+{
+    #[inline]
+    fn negative_to_usize(self, len: usize) -> Option<usize> {
+        if self >= Zero::zero() && (self.to_usize().unwrap()) < len {
+            Some(self.to_usize().unwrap())
         } else {
-            let subtract = self.unsigned_abs() as usize;
-            if subtract > index {
+            let subtract = self.abs().to_usize().unwrap();
+            if subtract > len {
                 None
             } else {
-                Some(index - subtract)
+                Some(len - subtract)
             }
         }
     }
diff --git a/polars/polars-lazy/Cargo.toml b/polars/polars-lazy/Cargo.toml
@@ -64,6 +64,7 @@ date_offset = ["polars-plan/date_offset"]
 trigonometry = ["polars-plan/trigonometry"]
 sign = ["polars-plan/sign"]
 timezones = ["polars-plan/timezones"]
+list_take = ["polars-ops/list_take", "polars-plan/list_take"]
 
 true_div = ["polars-plan/true_div"]
 
diff --git a/polars/polars-lazy/polars-plan/Cargo.toml b/polars/polars-lazy/polars-plan/Cargo.toml
@@ -52,6 +52,7 @@ dtype-struct = ["polars-core/dtype-struct"]
 dtype-binary = ["polars-core/dtype-binary"]
 object = ["polars-core/object"]
 date_offset = ["polars-time"]
+list_take = ["polars-ops/list_take"]
 trigonometry = []
 sign = []
 timezones = ["polars-time/timezones", "polars-core/timezones"]
diff --git a/polars/polars-lazy/polars-plan/src/dsl/function_expr/list.rs b/polars/polars-lazy/polars-plan/src/dsl/function_expr/list.rs
@@ -10,6 +10,8 @@ pub enum ListFunction {
     Contains,
     Slice,
     Get,
+    #[cfg(feature = "list_take")]
+    Take,
 }
 
 impl Display for ListFunction {
@@ -22,6 +24,8 @@ impl Display for ListFunction {
             Contains => "contains",
             Slice => "slice",
             Get => "get",
+            #[cfg(feature = "list_take")]
+            Take => "take",
         };
         write!(f, "{name}")
     }
@@ -185,3 +189,20 @@ pub(super) fn get(s: &mut [Series]) -> PolarsResult<Series> {
 
     }
 }
+
+#[cfg(feature = "list_take")]
+pub(super) fn take(args: &[Series]) -> PolarsResult<Series> {
+    let ca = &args[0];
+    let idx = &args[1];
+    let ca = ca.list()?;
+
+    if idx.len() == 1 {
+        // fast path
+        let idx = idx.get(0)?.try_extract::<i64>()?;
+        let out = ca.lst_get(idx)?;
+        // make sure we return a list
+        out.reshape(&[-1, 1])
+    } else {
+        ca.lst_take(idx)
+    }
+}
diff --git a/polars/polars-lazy/polars-plan/src/dsl/function_expr/mod.rs b/polars/polars-lazy/polars-plan/src/dsl/function_expr/mod.rs
@@ -326,6 +326,8 @@ impl From<FunctionExpr> for SpecialEq<Arc<dyn SeriesUdf>> {
                     Contains => wrap!(list::contains),
                     Slice => wrap!(list::slice),
                     Get => wrap!(list::get),
+                    #[cfg(feature = "list_take")]
+                    Take => map_as_slice!(list::take),
                 }
             }
             #[cfg(feature = "dtype-struct")]
diff --git a/polars/polars-lazy/polars-plan/src/dsl/function_expr/schema.rs b/polars/polars-lazy/polars-plan/src/dsl/function_expr/schema.rs
@@ -184,6 +184,8 @@ impl FunctionExpr {
                     Contains => with_dtype(DataType::Boolean),
                     Slice => same_type(),
                     Get => inner_type_list(),
+                    #[cfg(feature = "list_take")]
+                    Take => same_type(),
                 }
             }
             #[cfg(feature = "dtype-struct")]
diff --git a/polars/polars-lazy/polars-plan/src/dsl/list.rs b/polars/polars-lazy/polars-plan/src/dsl/list.rs
@@ -122,6 +122,14 @@ impl ListNameSpace {
             .map_many_private(FunctionExpr::ListExpr(ListFunction::Get), &[index], false)
     }
 
+    /// Get items in every sublist by multiple indexes.
+    #[cfg(feature = "list_take")]
+    #[cfg_attr(docsrs, doc(cfg(feature = "list_take")))]
+    pub fn take(self, index: Expr) -> Expr {
+        self.0
+            .map_many_private(FunctionExpr::ListExpr(ListFunction::Take), &[index], false)
+    }
+
     /// Get first item of every sublist.
     pub fn first(self) -> Expr {
         self.get(lit(0i64))
diff --git a/polars/polars-ops/Cargo.toml b/polars/polars-ops/Cargo.toml
@@ -35,6 +35,7 @@ dtype-i16 = ["polars-core/dtype-i16"]
 object = ["polars-core/object"]
 propagate_nans = []
 performant = ["polars-core/performant"]
+big_idx = ["polars-core/bigidx"]
 
 # ops
 to_dummies = []
@@ -56,3 +57,4 @@ cross_join = ["polars-core/cross_join"]
 chunked_ids = ["polars-core/chunked_ids"]
 asof_join = ["polars-core/asof_join"]
 semi_anti_join = ["polars-core/semi_anti_join"]
+list_take = []
diff --git a/polars/polars-ops/src/chunked_array/list/namespace.rs b/polars/polars-ops/src/chunked_array/list/namespace.rs
@@ -4,6 +4,8 @@ use std::fmt::Write;
 use polars_arrow::kernels::list::sublist_get;
 use polars_arrow::prelude::ValueSize;
 use polars_core::chunked_array::builder::get_list_builder;
+#[cfg(feature = "list_take")]
+use polars_core::export::num::{NumCast, Signed, Zero};
 #[cfg(feature = "diff")]
 use polars_core::series::ops::NullBehavior;
 use polars_core::utils::{try_get_supertype, CustomIterTools};
@@ -213,6 +215,72 @@ pub trait ListNameSpaceImpl: AsList {
         Series::try_from((ca.name(), chunks))
     }
 
+    #[cfg(feature = "list_take")]
+    fn lst_take(&self, idx: &Series) -> PolarsResult<Series> {
+        let list_ca = self.as_list();
+
+        let index_typed_index = |idx: &Series| {
+            let other = idx.cast(&IDX_DTYPE).unwrap();
+            let idx = other.idx().unwrap();
+            list_ca
+                .amortized_iter()
+                .map(|s| s.map(|s| s.as_ref().take(idx)).transpose())
+                .collect::<PolarsResult<ListChunked>>()
+                .map(|mut ca| {
+                    ca.rename(list_ca.name());
+                    ca.into_series()
+                })
+        };
+
+        use DataType::*;
+        match idx.dtype() {
+            List(_) => {
+                let idx_ca = idx.list().unwrap();
+                let mut out = list_ca
+                    .amortized_iter()
+                    .zip(idx_ca.into_iter())
+                    .map(|(opt_s, opt_idx)| {
+                        {
+                            match (opt_s, opt_idx) {
+                                (Some(s), Some(idx)) => take_series(s.as_ref(), idx),
+                                _ => None,
+                            }
+                        }
+                        .transpose()
+                    })
+                    .collect::<PolarsResult<ListChunked>>()?;
+                out.rename(list_ca.name());
+
+                Ok(out.into_series())
+            }
+            UInt32 | UInt64 => index_typed_index(idx),
+            dt if dt.is_signed() => {
+                if let Some(min) = idx.min::<i64>() {
+                    if min > 0 {
+                        let idx = idx.cast(&IDX_DTYPE).unwrap();
+                        index_typed_index(&idx)
+                    } else {
+                        let mut out = list_ca
+                            .amortized_iter()
+                            .map(|opt_s| {
+                                opt_s
+                                    .and_then(|s| take_series(s.as_ref(), idx.clone()))
+                                    .transpose()
+                            })
+                            .collect::<PolarsResult<ListChunked>>()?;
+                        out.rename(list_ca.name());
+                        Ok(out.into_series())
+                    }
+                } else {
+                    Err(PolarsError::ComputeError("All indices are null".into()))
+                }
+            }
+            dt => Err(PolarsError::ComputeError(
+                format!("Cannot use dtype: '{dt}' as index.").into(),
+            )),
+        }
+    }
+
     fn lst_concat(&self, other: &[Series]) -> PolarsResult<ListChunked> {
         let ca = self.as_list();
         let other_len = other.len();
@@ -360,3 +428,57 @@ pub trait ListNameSpaceImpl: AsList {
 }
 
 impl ListNameSpaceImpl for ListChunked {}
+
+#[cfg(feature = "list_take")]
+fn take_series(s: &Series, idx: Series) -> Option<PolarsResult<Series>> {
+    let len = s.len();
+    let idx = cast_index(idx, len);
+    let idx = idx.idx().unwrap();
+    Some(s.take(idx))
+}
+
+#[cfg(feature = "list_take")]
+fn cast_index_ca<T: PolarsNumericType>(idx: &ChunkedArray<T>, len: usize) -> Series
+where
+    T::Native: Copy + PartialOrd + PartialEq + NumCast + Signed + Zero,
+{
+    idx.into_iter()
+        .map(|opt_idx| opt_idx.and_then(|idx| idx.negative_to_usize(len).map(|idx| idx as IdxSize)))
+        .collect::<IdxCa>()
+        .into_series()
+}
+
+#[cfg(feature = "list_take")]
+fn cast_index(idx: Series, len: usize) -> Series {
+    use DataType::*;
+    match idx.dtype() {
+        #[cfg(feature = "big_idx")]
+        UInt32 => idx.cast(&IDX_DTYPE).unwrap(),
+        #[cfg(feature = "big_idx")]
+        UInt64 => idx,
+        #[cfg(not(feature = "big_idx"))]
+        UInt64 => idx.cast(&IDX_DTYPE).unwrap(),
+        #[cfg(not(feature = "big_idx"))]
+        UInt32 => idx,
+        dt if dt.is_unsigned() => idx.cast(&IDX_DTYPE).unwrap(),
+        Int8 => {
+            let a = idx.i8().unwrap();
+            cast_index_ca(a, len)
+        }
+        Int16 => {
+            let a = idx.i16().unwrap();
+            cast_index_ca(a, len)
+        }
+        Int32 => {
+            let a = idx.i32().unwrap();
+            cast_index_ca(a, len)
+        }
+        Int64 => {
+            let a = idx.i64().unwrap();
+            cast_index_ca(a, len)
+        }
+        _ => {
+            unreachable!()
+        }
+    }
+}
diff --git a/polars/src/lib.rs b/polars/src/lib.rs
@@ -232,6 +232,7 @@
 //!     - `interpolate` [interpolate None values](crate::chunked_array::ops::Interpolate)
 //!     - `extract_jsonpath` - [Run jsonpath queries on Utf8Chunked](https://goessner.net/articles/JsonPath/)
 //!     - `list` - List utils.
+//!         - `list_take` take sublist by multiple indices
 //!     - `rank` - Ranking algorithms.
 //!     - `moment` - kurtosis and skew statistics
 //!     - `ewma` - Exponential moving average windows
diff --git a/py-polars/Cargo.toml b/py-polars/Cargo.toml
@@ -67,6 +67,7 @@ performant = ["polars/performant"]
 timezones = ["polars/timezones"]
 cse = ["polars/cse"]
 merge_sorted = ["polars/merge_sorted"]
+list_take = ["polars/list_take"]
 
 all = [
   "json",
@@ -100,6 +101,7 @@ all = [
   "polars/binary_encoding",
   "streaming",
   "performant",
+  "list_take",
 ]
 
 # we cannot conditionaly activate simd
diff --git a/py-polars/docs/source/reference/expressions/array.rst b/py-polars/docs/source/reference/expressions/array.rst
@@ -30,5 +30,6 @@ The following methods are available under the `expr.arr` attribute.
     Expr.arr.sort
     Expr.arr.sum
     Expr.arr.tail
+    Expr.arr.take
     Expr.arr.to_struct
     Expr.arr.unique
diff --git a/py-polars/docs/source/reference/series/array.rst b/py-polars/docs/source/reference/series/array.rst
@@ -30,5 +30,6 @@ The following methods are available under the `Series.arr` attribute.
     Series.arr.sort
     Series.arr.sum
     Series.arr.tail
+    Series.arr.take
     Series.arr.to_struct
     Series.arr.unique
diff --git a/py-polars/polars/internals/expr/list.py b/py-polars/polars/internals/expr/list.py
@@ -274,6 +274,24 @@ def get(self, index: int | pli.Expr | str) -> pli.Expr:
         index = pli.expr_to_lit_or_expr(index, str_to_lit=False)._pyexpr
         return pli.wrap_expr(self._pyexpr.lst_get(index))
 
+    def take(self, index: pli.Expr | pli.Series | list[int]) -> pli.Expr:
+        """
+        Take sublists by multiple indices.
+
+        The indices may be defined in a single column, or by sublists in another
+        column of dtype ``List``.
+
+        Parameters
+        ----------
+        index
+            Indices to return per sublist
+
+        """
+        if isinstance(index, list):
+            index = pli.Series(index)
+        index = pli.expr_to_lit_or_expr(index, str_to_lit=False)._pyexpr
+        return pli.wrap_expr(self._pyexpr.lst_take(index))
+
     def __getitem__(self, item: int) -> pli.Expr:
         return self.get(item)
 
diff --git a/py-polars/polars/internals/series/list.py b/py-polars/polars/internals/series/list.py
@@ -84,6 +84,20 @@ def get(self, index: int | pli.Series | list[int]) -> pli.Series:
 
         """
 
+    def take(self, index: pli.Series | list[int]) -> pli.Series:
+        """
+        Take sublists by multiple indices.
+
+        The indices may be defined in a single column, or by sublists in another
+        column of dtype ``List``.
+
+        Parameters
+        ----------
+        index
+            Indices to return per sublist
+
+        """
+
     def __getitem__(self, item: int) -> pli.Series:
         return self.get(item)
 
diff --git a/py-polars/src/lazy/dsl.rs b/py-polars/src/lazy/dsl.rs
@@ -1458,6 +1458,11 @@ impl PyExpr {
         self.inner.clone().arr().get(index.inner).into()
     }
 
+    #[cfg(feature = "list_take")]
+    fn lst_take(&self, index: PyExpr) -> Self {
+        self.inner.clone().arr().take(index.inner).into()
+    }
+
     fn lst_join(&self, separator: &str) -> Self {
         self.inner.clone().arr().join(separator).into()
     }
diff --git a/py-polars/tests/unit/test_lists.py b/py-polars/tests/unit/test_lists.py

Original file line number	Diff line number	Diff line change
`@@ -326,6 +326,8 @@ impl From<FunctionExpr> for SpecialEq<Arc<dyn SeriesUdf>> {`
`326`	`326`	`Contains => wrap!(list::contains),`
`327`	`327`	`Slice => wrap!(list::slice),`
`328`	`328`	`Get => wrap!(list::get),`
	`329`	`+ #[cfg(feature = "list_take")]`
	`330`	`+ Take => map_as_slice!(list::take),`
`329`	`331`	`}`
`330`	`332`	`}`
`331`	`333`	`#[cfg(feature = "dtype-struct")]`
Original file line number	Diff line number	Diff line change
`@@ -184,6 +184,8 @@ impl FunctionExpr {`
`184`	`184`	`Contains => with_dtype(DataType::Boolean),`
`185`	`185`	`Slice => same_type(),`
`186`	`186`	`Get => inner_type_list(),`
	`187`	`+ #[cfg(feature = "list_take")]`
	`188`	`+ Take => same_type(),`
`187`	`189`	`}`
`188`	`190`	`}`
`189`	`191`	`#[cfg(feature = "dtype-struct")]`