polars-df 0.21.0 → 0.22.0

data/ext/polars/src/series/general.rs

@@ -1,9 +1,12 @@
-use magnus::{Error, IntoValue, Value, exception};
+use magnus::{Error, IntoValue, RArray, Ruby, Value, value::ReprValue};
 use polars::prelude::*;
 use polars::series::IsSorted;
+use polars_core::utils::flatten::flatten_series;
 
 use crate::conversion::*;
-use crate::{RbDataFrame, RbPolarsErr, RbResult, RbSeries};
+use crate::exceptions::RbIndexError;
+use crate::rb_modules;
+use crate::{RbDataFrame, RbErr, RbPolarsErr, RbResult, RbSeries};
 
 impl RbSeries {
     pub fn struct_unnest(&self) -> RbResult<RbDataFrame> {
@@ -13,7 +16,6 @@ impl RbSeries {
         Ok(df.into())
     }
 
-    // TODO add to Ruby
     pub fn struct_fields(&self) -> RbResult<Vec<String>> {
         let binding = self.series.borrow();
         let ca = binding.struct_().map_err(RbPolarsErr::from)?;
@@ -84,8 +86,39 @@ impl RbSeries {
         }
     }
 
-    pub fn get_idx(&self, idx: usize) -> RbResult<Value> {
-        Ok(Wrap(self.series.borrow().get(idx).map_err(RbPolarsErr::from)?).into_value())
+    pub fn get_index(ruby: &Ruby, rb_self: &Self, index: usize) -> RbResult<Value> {
+        let binding = rb_self.series.borrow();
+        let av = match binding.get(index) {
+            Ok(v) => v,
+            Err(PolarsError::OutOfBounds(err)) => {
+                return Err(RbIndexError::new_err(err.to_string()));
+            }
+            Err(e) => return Err(RbPolarsErr::from(e).into()),
+        };
+
+        match av {
+            AnyValue::List(s) | AnyValue::Array(s, _) => {
+                let rbseries = RbSeries::new(s);
+                rb_modules::pl_utils().funcall("wrap_s", (rbseries,))
+            }
+            _ => Ok(Wrap(av).into_value_with(ruby)),
+        }
+    }
+
+    pub fn get_index_signed(ruby: &Ruby, rb_self: &Self, index: isize) -> RbResult<Value> {
+        let index = if index < 0 {
+            match rb_self.len().checked_sub(index.unsigned_abs()) {
+                Some(v) => v,
+                None => {
+                    return Err(RbIndexError::new_err(
+                        polars_err!(oob = index, rb_self.len()).to_string(),
+                    ));
+                }
+            }
+        } else {
+            usize::try_from(index).unwrap()
+        };
+        Self::get_index(ruby, rb_self, index)
     }
 
     pub fn bitand(&self, other: &RbSeries) -> RbResult<Self> {
@@ -115,16 +148,17 @@ impl RbSeries {
         self.series.borrow_mut().rename(name.into());
     }
 
-    pub fn dtype(&self) -> Value {
-        Wrap(self.series.borrow().dtype().clone()).into_value()
+    pub fn dtype(ruby: &Ruby, rb_self: &Self) -> Value {
+        Wrap(rb_self.series.borrow().dtype().clone()).into_value_with(ruby)
     }
 
-    pub fn inner_dtype(&self) -> Option<Value> {
-        self.series
+    pub fn inner_dtype(ruby: &Ruby, rb_self: &Self) -> Option<Value> {
+        rb_self
+            .series
             .borrow()
             .dtype()
             .inner_dtype()
-            .map(|dt| Wrap(dt.clone()).into_value())
+            .map(|dt| Wrap(dt.clone()).into_value_with(ruby))
     }
 
     pub fn set_sorted_flag(&self, descending: bool) -> Self {
@@ -141,11 +175,11 @@ impl RbSeries {
         self.series.borrow().n_chunks()
     }
 
-    pub fn append(&self, other: &RbSeries) -> RbResult<()> {
-        let mut binding = self.series.borrow_mut();
+    pub fn append(ruby: &Ruby, rb_self: &Self, other: &RbSeries) -> RbResult<()> {
+        let mut binding = rb_self.series.borrow_mut();
         let res = binding.append(&other.series.borrow());
         if let Err(e) = res {
-            Err(Error::new(exception::runtime_error(), e.to_string()))
+            Err(Error::new(ruby.exception_runtime_error(), e.to_string()))
         } else {
             Ok(())
         }
@@ -159,22 +193,30 @@ impl RbSeries {
         Ok(())
     }
 
-    pub fn new_from_index(&self, index: usize, length: usize) -> RbResult<Self> {
-        if index >= self.series.borrow().len() {
-            Err(Error::new(exception::arg_error(), "index is out of bounds"))
+    pub fn new_from_index(
+        ruby: &Ruby,
+        rb_self: &Self,
+        index: usize,
+        length: usize,
+    ) -> RbResult<Self> {
+        if index >= rb_self.series.borrow().len() {
+            Err(Error::new(
+                ruby.exception_arg_error(),
+                "index is out of bounds",
+            ))
         } else {
-            Ok(self.series.borrow().new_from_index(index, length).into())
+            Ok(rb_self.series.borrow().new_from_index(index, length).into())
         }
     }
 
-    pub fn filter(&self, filter: &RbSeries) -> RbResult<Self> {
+    pub fn filter(ruby: &Ruby, rb_self: &Self, filter: &RbSeries) -> RbResult<Self> {
         let filter_series = &filter.series.borrow();
         if let Ok(ca) = filter_series.bool() {
-            let series = self.series.borrow().filter(ca).unwrap();
+            let series = rb_self.series.borrow().filter(ca).unwrap();
             Ok(series.into())
         } else {
             Err(Error::new(
-                exception::runtime_error(),
+                ruby.exception_runtime_error(),
                 "Expected a boolean mask".to_string(),
             ))
         }
@@ -279,12 +321,42 @@ impl RbSeries {
         }
     }
 
-    pub fn not(&self) -> RbResult<Self> {
+    pub fn not_(&self) -> RbResult<Self> {
         let binding = self.series.borrow();
         let bool = binding.bool().map_err(RbPolarsErr::from)?;
         Ok((!bool).into_series().into())
     }
 
+    pub fn shrink_dtype(&self) -> RbResult<Self> {
+        self.series
+            .borrow()
+            .shrink_type()
+            .map(Into::into)
+            .map_err(RbPolarsErr::from)
+            .map_err(RbErr::from)
+    }
+
+    pub fn str_to_decimal_infer(&self, inference_length: usize) -> RbResult<Self> {
+        let s = self.series.borrow();
+        let ca = s.str().map_err(RbPolarsErr::from)?;
+        ca.to_decimal_infer(inference_length)
+            .map(Series::from)
+            .map(Into::into)
+            .map_err(RbPolarsErr::from)
+            .map_err(RbErr::from)
+    }
+
+    pub fn str_json_decode(&self, infer_schema_length: Option<usize>) -> RbResult<Self> {
+        let lock = self.series.borrow();
+        lock.str()
+            .map_err(RbPolarsErr::from)?
+            .json_decode(None, infer_schema_length)
+            .map(|s| s.with_name(lock.name().clone()))
+            .map(Into::into)
+            .map_err(RbPolarsErr::from)
+            .map_err(RbErr::from)
+    }
+
     pub fn to_s(&self) -> String {
         format!("{}", self.series.borrow())
     }
@@ -370,6 +442,33 @@ impl RbSeries {
         Ok(out.into())
     }
 
+    pub fn get_chunks(ruby: &Ruby, rb_self: &Self) -> RbResult<RArray> {
+        ruby.ary_try_from_iter(
+            flatten_series(&rb_self.series.borrow())
+                .into_iter()
+                .map(|s| rb_modules::pl_utils().funcall::<_, _, Value>("wrap_s", (Self::new(s),))),
+        )
+    }
+
+    pub fn is_sorted(&self, descending: bool, nulls_last: bool) -> RbResult<bool> {
+        let options = SortOptions {
+            descending,
+            nulls_last,
+            multithreaded: true,
+            maintain_order: false,
+            limit: None,
+        };
+        Ok(self
+            .series
+            .borrow()
+            .is_sorted(options)
+            .map_err(RbPolarsErr::from)?)
+    }
+
+    pub fn clear(&self) -> Self {
+        self.series.borrow().clear().into()
+    }
+
     pub fn time_unit(&self) -> Option<String> {
         if let DataType::Datetime(tu, _) | DataType::Duration(tu) = self.series.borrow().dtype() {
             Some(

data/ext/polars/src/series/mod.rs

@@ -8,13 +8,14 @@ mod import;
 mod map;
 mod scatter;
 
-use magnus::{RArray, prelude::*};
+use magnus::{DataTypeFunctions, RArray, Ruby, TypedData, gc, prelude::*};
 use polars::prelude::*;
 use std::cell::RefCell;
 
-use crate::RbResult;
+use crate::{ObjectValue, RbResult};
 
-#[magnus::wrap(class = "Polars::RbSeries")]
+#[derive(TypedData)]
+#[magnus(class = "Polars::RbSeries", mark)]
 pub struct RbSeries {
     pub series: RefCell<Series>,
 }
@@ -42,9 +43,33 @@ pub fn to_series(rs: RArray) -> RbResult<Vec<Series>> {
 }
 
 pub fn to_rbseries(s: Vec<Column>) -> RArray {
-    RArray::from_iter(
+    Ruby::get().unwrap().ary_from_iter(
         s.into_iter()
             .map(|c| c.take_materialized_series())
             .map(RbSeries::new),
     )
 }
+
+pub fn mark_series(marker: &gc::Marker, series: &Series) {
+    if let DataType::Object(_) = series.dtype() {
+        for i in 0..series.len() {
+            let obj: Option<&ObjectValue> = series.get_object(i).map(|any| any.into());
+            if let Some(o) = obj {
+                marker.mark(o.inner);
+            }
+        }
+    }
+}
+
+impl DataTypeFunctions for RbSeries {
+    fn mark(&self, marker: &gc::Marker) {
+        // this is not ideal, as objects will not be marked if unable to borrow
+        // this should never happen, but log for now to avoid panic,
+        // as most series will not use Object datatype
+        if let Ok(s) = &self.series.try_borrow() {
+            mark_series(marker, s);
+        } else {
+            eprintln!("[polars] Could not borrow!");
+        }
+    }
+}

data/lib/polars/array_expr.rb

@@ -9,6 +9,181 @@ module Polars
       self._rbexpr = expr._rbexpr
     end
 
+    # Return the number of elements in each array.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2], [4, 3]]},
+    #     schema: {"a" => Polars::Array.new(Polars::Int64, 2)}
+    #   )
+    #   df.select(Polars.col("a").arr.len)
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌─────┐
+    #   # │ a   │
+    #   # │ --- │
+    #   # │ u32 │
+    #   # ╞═════╡
+    #   # │ 2   │
+    #   # │ 2   │
+    #   # └─────┘
+    def len
+      Utils.wrap_expr(_rbexpr.arr_len)
+    end
+
+    # Slice every subarray.
+    #
+    # @param offset [Integer]
+    #   Start index. Negative indexing is supported.
+    # @param length [Integer]
+    #   Length of the slice. If set to `nil` (default), the slice is taken to the
+    #   end of the list.
+    # @param as_array [Boolean]
+    #   Return result as a fixed-length `Array`, otherwise as a `List`.
+    #   If true `length` and `offset` must be constant values.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2], [4, 3]]},
+    #     schema: {"a" => Polars::Array.new(Polars::Int64, 2)}
+    #   )
+    #   df.select(Polars.col("a").arr.slice(0, 1))
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌───────────┐
+    #   # │ a         │
+    #   # │ ---       │
+    #   # │ list[i64] │
+    #   # ╞═══════════╡
+    #   # │ [1]       │
+    #   # │ [4]       │
+    #   # └───────────┘
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2], [4, 3]]},
+    #     schema: {"a" => Polars::Array.new(Polars::Int64, 2)}
+    #   )
+    #   df.select(Polars.col("a").arr.slice(0, 1, as_array: true))
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌───────────────┐
+    #   # │ a             │
+    #   # │ ---           │
+    #   # │ array[i64, 1] │
+    #   # ╞═══════════════╡
+    #   # │ [1]           │
+    #   # │ [4]           │
+    #   # └───────────────┘
+    def slice(
+      offset,
+      length = nil,
+      as_array: false
+    )
+      offset = Utils.parse_into_expression(offset)
+      length = !length.nil? ? Utils.parse_into_expression(length) : nil
+      Utils.wrap_expr(_rbexpr.arr_slice(offset, length, as_array))
+    end
+
+    # Get the first `n` elements of the sub-arrays.
+    #
+    # @param n [Integer]
+    #   Number of values to return for each sublist.
+    # @param as_array [Boolean]
+    #   Return result as a fixed-length `Array`, otherwise as a `List`.
+    #   If true `n` must be a constant value.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2], [4, 3]]},
+    #     schema: {"a" => Polars::Array.new(Polars::Int64, 2)}
+    #   )
+    #   df.select(Polars.col("a").arr.head(1))
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌───────────┐
+    #   # │ a         │
+    #   # │ ---       │
+    #   # │ list[i64] │
+    #   # ╞═══════════╡
+    #   # │ [1]       │
+    #   # │ [4]       │
+    #   # └───────────┘
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2], [4, 3]]},
+    #     schema: {"a" => Polars::Array.new(Polars::Int64, 2)}
+    #   )
+    #   df.select(Polars.col("a").arr.head(1, as_array: true))
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌───────────────┐
+    #   # │ a             │
+    #   # │ ---           │
+    #   # │ array[i64, 1] │
+    #   # ╞═══════════════╡
+    #   # │ [1]           │
+    #   # │ [4]           │
+    #   # └───────────────┘
+    def head(n = 5, as_array: false)
+      slice(0, n, as_array: as_array)
+    end
+
+    # Slice the last `n` values of every sublist.
+    #
+    # @param n [Integer]
+    #   Number of values to return for each sublist.
+    # @param as_array [Boolean]
+    #   Return result as a fixed-length `Array`, otherwise as a `List`.
+    #   If true `n` must be a constant value.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2], [4, 3]]},
+    #     schema: {"a" => Polars::Array.new(Polars::Int64, 2)}
+    #   )
+    #   df.select(Polars.col("a").arr.tail(1))
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌───────────┐
+    #   # │ a         │
+    #   # │ ---       │
+    #   # │ list[i64] │
+    #   # ╞═══════════╡
+    #   # │ [2]       │
+    #   # │ [3]       │
+    #   # └───────────┘
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2], [4, 3]]},
+    #     schema: {"a" => Polars::Array.new(Polars::Int64, 2)}
+    #   )
+    #   df.select(Polars.col("a").arr.tail(1, as_array: true))
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌───────────────┐
+    #   # │ a             │
+    #   # │ ---           │
+    #   # │ array[i64, 1] │
+    #   # ╞═══════════════╡
+    #   # │ [2]           │
+    #   # │ [3]           │
+    #   # └───────────────┘
+    def tail(n = 5, as_array: false)
+      n = Utils.parse_into_expression(n)
+      Utils.wrap_expr(_rbexpr.arr_tail(n, as_array))
+    end
+
     # Compute the min values of the sub-arrays.
     #
     # @return [Expr]
@@ -30,7 +205,7 @@ module Polars
     #   # │ 3   │
     #   # └─────┘
     def min
-      Utils.wrap_expr(_rbexpr.array_min)
+      Utils.wrap_expr(_rbexpr.arr_min)
     end
 
     # Compute the max values of the sub-arrays.
@@ -54,7 +229,7 @@ module Polars
     #   # │ 4   │
     #   # └─────┘
     def max
-      Utils.wrap_expr(_rbexpr.array_max)
+      Utils.wrap_expr(_rbexpr.arr_max)
     end
 
     # Compute the sum values of the sub-arrays.
@@ -78,7 +253,103 @@ module Polars
     #   # │ 7   │
     #   # └─────┘
     def sum
-      Utils.wrap_expr(_rbexpr.array_sum)
+      Utils.wrap_expr(_rbexpr.arr_sum)
+    end
+
+    # Compute the std of the values of the sub-arrays.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2], [4, 3]]},
+    #     schema: {"a" => Polars::Array.new(2, Polars::Int64)}
+    #   )
+    #   df.select(Polars.col("a").arr.std)
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌──────────┐
+    #   # │ a        │
+    #   # │ ---      │
+    #   # │ f64      │
+    #   # ╞══════════╡
+    #   # │ 0.707107 │
+    #   # │ 0.707107 │
+    #   # └──────────┘
+    def std(ddof: 1)
+      Utils.wrap_expr(_rbexpr.arr_std(ddof))
+    end
+
+    # Compute the var of the values of the sub-arrays.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2], [4, 3]]},
+    #     schema: {"a" => Polars::Array.new(2, Polars::Int64)}
+    #   )
+    #   df.select(Polars.col("a").arr.var)
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌─────┐
+    #   # │ a   │
+    #   # │ --- │
+    #   # │ f64 │
+    #   # ╞═════╡
+    #   # │ 0.5 │
+    #   # │ 0.5 │
+    #   # └─────┘
+    def var(ddof: 1)
+      Utils.wrap_expr(_rbexpr.arr_var(ddof))
+    end
+
+    # Compute the mean of the values of the sub-arrays.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2, 3], [1, 1, 16]]},
+    #     schema: {"a" => Polars::Array.new(Polars::Int64, 3)}
+    #   )
+    #   df.select(Polars.col("a").arr.mean)
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌─────┐
+    #   # │ a   │
+    #   # │ --- │
+    #   # │ f64 │
+    #   # ╞═════╡
+    #   # │ 2.0 │
+    #   # │ 6.0 │
+    #   # └─────┘
+    def mean
+      Utils.wrap_expr(_rbexpr.arr_mean)
+    end
+
+    # Compute the median of the values of the sub-arrays.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2], [4, 3]]},
+    #     schema: {"a" => Polars::Array.new(2, Polars::Int64)}
+    #   )
+    #   df.select(Polars.col("a").arr.median)
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌─────┐
+    #   # │ a   │
+    #   # │ --- │
+    #   # │ f64 │
+    #   # ╞═════╡
+    #   # │ 1.5 │
+    #   # │ 3.5 │
+    #   # └─────┘
+    def median
+      Utils.wrap_expr(_rbexpr.arr_median)
     end
 
     # Get the unique/distinct values in the array.
@@ -109,6 +380,32 @@ module Polars
       Utils.wrap_expr(_rbexpr.arr_unique(maintain_order))
     end
 
+    # Count the number of unique values in every sub-arrays.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "a" => [[1, 1, 2], [2, 3, 4]],
+    #     },
+    #     schema: {"a" => Polars::Array.new(Polars::Int64, 3)}
+    #   )
+    #   df.with_columns(n_unique: Polars.col("a").arr.n_unique)
+    #   # =>
+    #   # shape: (2, 2)
+    #   # ┌───────────────┬──────────┐
+    #   # │ a             ┆ n_unique │
+    #   # │ ---           ┆ ---      │
+    #   # │ array[i64, 3] ┆ u32      │
+    #   # ╞═══════════════╪══════════╡
+    #   # │ [1, 1, 2]     ┆ 2        │
+    #   # │ [2, 3, 4]     ┆ 3        │
+    #   # └───────────────┴──────────┘
+    def n_unique
+      Utils.wrap_expr(_rbexpr.arr_n_unique)
+    end
+
     # Convert an Array column into a List column with the same inner data type.
     #
     # @return [Expr]
@@ -535,5 +832,87 @@ module Polars
       element = Utils.parse_into_expression(element, str_as_lit: true)
       Utils.wrap_expr(_rbexpr.arr_count_matches(element))
     end
+
+    # Convert the Series of type `Array` to a Series of type `Struct`.
+    #
+    # @param fields [Object]
+    #   If the name and number of the desired fields is known in advance
+    #   a list of field names can be given, which will be assigned by index.
+    #   Otherwise, to dynamically assign field names, a custom function can be
+    #   used; if neither are set, fields will be `field_0, field_1 .. field_n`.
+    #
+    # @return [Expr]
+    #
+    # @example Convert array to struct with default field name assignment:
+    #   df = Polars::DataFrame.new(
+    #     {"n" => [[0, 1, 2], [3, 4, 5]]}, schema: {"n" => Polars::Array.new(Polars::Int8, 3)}
+    #   )
+    #   df.with_columns(struct: Polars.col("n").arr.to_struct)
+    #   # =>
+    #   # shape: (2, 2)
+    #   # ┌──────────────┬───────────┐
+    #   # │ n            ┆ struct    │
+    #   # │ ---          ┆ ---       │
+    #   # │ array[i8, 3] ┆ struct[3] │
+    #   # ╞══════════════╪═══════════╡
+    #   # │ [0, 1, 2]    ┆ {0,1,2}   │
+    #   # │ [3, 4, 5]    ┆ {3,4,5}   │
+    #   # └──────────────┴───────────┘
+    def to_struct(fields: nil)
+      raise Todo if fields
+      if fields.is_a?(Enumerable)
+        field_names = fields.to_a
+        rbexpr = _rbexpr.arr_to_struct(nil)
+        Utils.wrap_expr(rbexpr).struct.rename_fields(field_names)
+      else
+        rbexpr = _rbexpr.arr_to_struct(fields)
+        Utils.wrap_expr(rbexpr)
+      end
+    end
+
+    # Shift array values by the given number of indices.
+    #
+    # @param n [Integer]
+    #   Number of indices to shift forward. If a negative value is passed, values
+    #   are shifted in the opposite direction instead.
+    #
+    # @return [Expr]
+    #
+    # @note
+    #   This method is similar to the `LAG` operation in SQL when the value for `n`
+    #   is positive. With a negative value for `n`, it is similar to `LEAD`.
+    #
+    # @example By default, array values are shifted forward by one index.
+    #   df = Polars::DataFrame.new(
+    #     {"a" => [[1, 2, 3], [4, 5, 6]]}, schema: {"a" => Polars::Array.new(Polars::Int64, 3)}
+    #   )
+    #   df.with_columns(shift: Polars.col("a").arr.shift)
+    #   # =>
+    #   # shape: (2, 2)
+    #   # ┌───────────────┬───────────────┐
+    #   # │ a             ┆ shift         │
+    #   # │ ---           ┆ ---           │
+    #   # │ array[i64, 3] ┆ array[i64, 3] │
+    #   # ╞═══════════════╪═══════════════╡
+    #   # │ [1, 2, 3]     ┆ [null, 1, 2]  │
+    #   # │ [4, 5, 6]     ┆ [null, 4, 5]  │
+    #   # └───────────────┴───────────────┘
+    #
+    # @example Pass a negative value to shift in the opposite direction instead.
+    #   df.with_columns(shift: Polars.col("a").arr.shift(-2))
+    #   # =>
+    #   # shape: (2, 2)
+    #   # ┌───────────────┬─────────────────┐
+    #   # │ a             ┆ shift           │
+    #   # │ ---           ┆ ---             │
+    #   # │ array[i64, 3] ┆ array[i64, 3]   │
+    #   # ╞═══════════════╪═════════════════╡
+    #   # │ [1, 2, 3]     ┆ [3, null, null] │
+    #   # │ [4, 5, 6]     ┆ [6, null, null] │
+    #   # └───────────────┴─────────────────┘
+    def shift(n = 1)
+      n = Utils.parse_into_expression(n)
+      Utils.wrap_expr(_rbexpr.arr_shift(n))
+    end
   end
 end