polars_ops/chunked_array/list/

namespace.rs

use std::fmt::Write;

use arrow::array::ValueSize;
use polars_compute::gather::sublist::list::{index_is_oob, sublist_get};
use polars_core::chunked_array::builder::get_list_builder;
#[cfg(feature = "diff")]
use polars_core::series::ops::NullBehavior;
use polars_core::utils::{CustomIterTools, try_get_supertype};

use super::*;
use crate::chunked_array::list::min_max::{list_max_function, list_min_function};
use crate::chunked_array::list::sum_mean::sum_with_nulls;
#[cfg(feature = "diff")]
use crate::prelude::diff;
use crate::prelude::list::sum_mean::{mean_list_numerical, sum_list_numerical};
use crate::series::{ArgAgg, convert_and_bound_index};

pub(super) fn has_inner_nulls(ca: &ListChunked) -> bool {
    for arr in ca.downcast_iter() {
        if arr.values().null_count() > 0 {
            return true;
        }
    }
    false
}

fn cast_rhs(
    other: &mut [Column],
    inner_type: &DataType,
    dtype: &DataType,
    length: usize,
    allow_broadcast: bool,
) -> PolarsResult<()> {
    for s in other.iter_mut() {
        // make sure that inner types match before we coerce into list
        if !matches!(s.dtype(), DataType::List(_)) {
            *s = s.cast(inner_type)?
        }
        if !matches!(s.dtype(), DataType::List(_)) && s.dtype() == inner_type {
            // coerce to list JIT
            *s = s
                .reshape_list(&[ReshapeDimension::Infer, ReshapeDimension::new_dimension(1)])
                .unwrap();
        }
        if s.dtype() != dtype {
            *s = s.cast(dtype).map_err(|e| {
                polars_err!(
                    SchemaMismatch:
                    "cannot concat `{}` into a list of `{}`: {}",
                    s.dtype(),
                    dtype,
                    e
                )
            })?;
        }

        if allow_broadcast {
            // broadcast JIT
            s.broadcast_in_place_to(length)?;
        } else {
            polars_ensure!(
                s.len() == length || s.len() == 1,
                ShapeMismatch: "series length {} does not match expected length of {}",
                s.len(), length
            );
        }
    }
    Ok(())
}

pub trait ListNameSpaceImpl: AsList {
    /// In case the inner dtype [`DataType::String`], the individual items will be joined into a
    /// single string separated by `separator`.
    fn lst_join(
        &self,
        separator: &StringChunked,
        ignore_nulls: bool,
    ) -> PolarsResult<StringChunked> {
        let ca = self.as_list();
        match ca.inner_dtype() {
            DataType::String => match separator.len() {
                1 => match separator.get(0) {
                    Some(separator) => self.join_literal(separator, ignore_nulls),
                    _ => Ok(StringChunked::full_null(ca.name().clone(), ca.len())),
                },
                _ => self.join_many(separator, ignore_nulls),
            },
            dt => polars_bail!(op = "`lst.join`", got = dt, expected = "String"),
        }
    }

    fn join_literal(&self, separator: &str, ignore_nulls: bool) -> PolarsResult<StringChunked> {
        let ca = self.as_list();
        // used to amortize heap allocs
        let mut buf = String::with_capacity(128);
        let mut builder = StringChunkedBuilder::new(ca.name().clone(), ca.len());

        ca.for_each_amortized(|opt_s| {
            let opt_val = opt_s.and_then(|s| {
                // make sure that we don't write values of previous iteration
                buf.clear();
                let ca = s.as_ref().str().unwrap();

                if ca.null_count() != 0 && !ignore_nulls {
                    return None;
                }

                for arr in ca.downcast_iter() {
                    for val in arr.non_null_values_iter() {
                        buf.write_str(val).unwrap();
                        buf.write_str(separator).unwrap();
                    }
                }

                // last value should not have a separator, so slice that off
                // saturating sub because there might have been nothing written.
                Some(&buf[..buf.len().saturating_sub(separator.len())])
            });
            builder.append_option(opt_val)
        });
        Ok(builder.finish())
    }

    fn join_many(
        &self,
        separator: &StringChunked,
        ignore_nulls: bool,
    ) -> PolarsResult<StringChunked> {
        let ca = self.as_list();
        // used to amortize heap allocs
        let mut buf = String::with_capacity(128);
        let mut builder = StringChunkedBuilder::new(ca.name().clone(), ca.len());
        {
            ca.amortized_iter()
                .zip(separator.iter())
                .for_each(|(opt_s, opt_sep)| match opt_sep {
                    Some(separator) => {
                        let opt_val = opt_s.and_then(|s| {
                            // make sure that we don't write values of previous iteration
                            buf.clear();
                            let ca = s.as_ref().str().unwrap();

                            if ca.null_count() != 0 && !ignore_nulls {
                                return None;
                            }

                            for arr in ca.downcast_iter() {
                                for val in arr.non_null_values_iter() {
                                    buf.write_str(val).unwrap();
                                    buf.write_str(separator).unwrap();
                                }
                            }

                            // last value should not have a separator, so slice that off
                            // saturating sub because there might have been nothing written.
                            Some(&buf[..buf.len().saturating_sub(separator.len())])
                        });
                        builder.append_option(opt_val)
                    },
                    _ => builder.append_null(),
                })
        }
        Ok(builder.finish())
    }

    fn lst_max(&self) -> PolarsResult<Series> {
        list_max_function(self.as_list())
    }

    fn lst_min(&self) -> PolarsResult<Series> {
        list_min_function(self.as_list())
    }

    fn lst_sum(&self) -> PolarsResult<Series> {
        let ca = self.as_list();

        if has_inner_nulls(ca) {
            return sum_with_nulls(ca, ca.inner_dtype());
        };

        match ca.inner_dtype() {
            DataType::Boolean => Ok(count_boolean_bits(ca).into_series()),
            dt if dt.is_primitive_numeric() => Ok(sum_list_numerical(ca, dt)),
            dt => sum_with_nulls(ca, dt),
        }
    }

    fn lst_mean(&self) -> Series {
        let ca = self.as_list();

        if has_inner_nulls(ca) {
            return sum_mean::mean_with_nulls(ca);
        };

        match ca.inner_dtype() {
            dt if dt.is_primitive_numeric() => mean_list_numerical(ca, dt),
            _ => sum_mean::mean_with_nulls(ca),
        }
    }

    fn lst_median(&self) -> Series {
        let ca = self.as_list();
        dispersion::median_with_nulls(ca)
    }

    fn lst_std(&self, ddof: u8) -> Series {
        let ca = self.as_list();
        dispersion::std_with_nulls(ca, ddof)
    }

    fn lst_var(&self, ddof: u8) -> PolarsResult<Series> {
        let ca = self.as_list();
        dispersion::var_with_nulls(ca, ddof)
    }

    fn same_type(&self, out: ListChunked) -> ListChunked {
        let ca = self.as_list();
        let dtype = ca.dtype();
        if out.dtype() != dtype {
            out.cast(ca.dtype()).unwrap().list().unwrap().clone()
        } else {
            out
        }
    }

    fn lst_sort(&self, options: SortOptions) -> PolarsResult<ListChunked> {
        let ca = self.as_list();
        // SAFETY: `sort_with`` doesn't change the dtype
        let out = unsafe { ca.try_apply_amortized_same_type(|s| s.as_ref().sort_with(options))? };
        Ok(self.same_type(out))
    }

    fn lst_arg_min(&self) -> IdxCa {
        let ca = self.as_list();
        ca.apply_amortized_generic(|opt_s| {
            opt_s.and_then(|s| s.as_ref().arg_min().map(|idx| idx as IdxSize))
        })
    }

    fn lst_arg_max(&self) -> IdxCa {
        let ca = self.as_list();
        ca.apply_amortized_generic(|opt_s| {
            opt_s.and_then(|s| s.as_ref().arg_max().map(|idx| idx as IdxSize))
        })
    }

    #[cfg(feature = "diff")]
    fn lst_diff(&self, n: i64, null_behavior: NullBehavior) -> PolarsResult<ListChunked> {
        let ca = self.as_list();
        ca.try_apply_amortized(|s| diff(s.as_ref(), n, null_behavior))
    }

    fn lst_shift(&self, periods: &Column) -> PolarsResult<ListChunked> {
        let ca = self.as_list();
        let periods_s = periods.cast(&DataType::Int64)?;
        let periods = periods_s.i64()?;

        polars_ensure!(
            ca.len() == periods.len() || ca.len() == 1 || periods.len() == 1,
            length_mismatch = "list.shift",
            ca.len(),
            periods.len()
        );

        let target_len = periods.len();
        if ca.len() == 1 && target_len > 1 {
            let single_list = ca.get_as_series(0);
            let out = shift_broadcast_list(
                single_list,
                periods,
                target_len,
                ca.name().clone(),
                ca.inner_dtype(),
            );
            return Ok(self.same_type(out));
        }

        let out = match periods.len() {
            1 => {
                if let Some(periods) = periods.get(0) {
                    // SAFETY: `shift` doesn't change the dtype
                    unsafe { ca.apply_amortized_same_type(|s| s.as_ref().shift(periods)) }
                } else {
                    ListChunked::full_null_with_dtype(ca.name().clone(), ca.len(), ca.inner_dtype())
                }
            },
            _ => ca.zip_and_apply_amortized(periods, |opt_s, opt_periods| {
                match (opt_s, opt_periods) {
                    (Some(s), Some(periods)) => Some(s.as_ref().shift(periods)),
                    _ => None,
                }
            }),
        };
        Ok(self.same_type(out))
    }

    fn lst_slice(&self, offset: i64, length: usize) -> ListChunked {
        let ca = self.as_list();
        // SAFETY: `slice` doesn't change the dtype
        unsafe { ca.apply_amortized_same_type(|s| s.as_ref().slice(offset, length)) }
    }

    fn lst_lengths(&self) -> IdxCa {
        let ca = self.as_list();

        let ca_validity = ca.rechunk_validity();

        if ca_validity.as_ref().is_some_and(|x| x.set_bits() == 0) {
            return IdxCa::full_null(ca.name().clone(), ca.len());
        }

        let mut lengths = Vec::with_capacity(ca.len());
        ca.downcast_iter().for_each(|arr| {
            let offsets = arr.offsets().as_slice();
            let mut last = offsets[0];
            for o in &offsets[1..] {
                lengths.push((*o - last) as IdxSize);
                last = *o;
            }
        });

        let arr = IdxArr::from_vec(lengths).with_validity(ca_validity);
        IdxCa::with_chunk(ca.name().clone(), arr)
    }

    /// Get the value by index in the sublists.
    /// So index `0` would return the first item of every sublist
    /// and index `-1` would return the last item of every sublist
    /// if an index is out of bounds, it will return a `None`.
    fn lst_get(&self, idx: i64, null_on_oob: bool) -> PolarsResult<Series> {
        let ca = self.as_list();
        if !null_on_oob && ca.downcast_iter().any(|arr| index_is_oob(arr, idx)) {
            polars_bail!(ComputeError: "get index is out of bounds");
        }

        let chunks = ca
            .downcast_iter()
            .map(|arr| sublist_get(arr, idx))
            .collect::<Vec<_>>();

        let s = Series::try_from((ca.name().clone(), chunks)).unwrap();
        // SAFETY: every element in list has dtype equal to its inner type
        unsafe { s.from_physical_unchecked(ca.inner_dtype()) }
    }

    #[cfg(feature = "list_gather")]
    fn lst_gather_every(&self, n: &IdxCa, offset: &IdxCa) -> PolarsResult<Series> {
        let list_ca = self.as_list();
        let out = match (n.len(), offset.len()) {
            (1, 1) => match (n.get(0), offset.get(0)) {
                (Some(n), Some(offset)) => unsafe {
                    // SAFETY: `gather_every` doesn't change the dtype
                    list_ca.try_apply_amortized_same_type(|s| {
                        s.as_ref().gather_every(n as usize, offset as usize)
                    })?
                },
                _ => ListChunked::full_null_with_dtype(
                    list_ca.name().clone(),
                    list_ca.len(),
                    list_ca.inner_dtype(),
                ),
            },
            (1, len_offset) if len_offset == list_ca.len() => {
                if let Some(n) = n.get(0) {
                    list_ca.try_zip_and_apply_amortized(offset, |opt_s, opt_offset| {
                        match (opt_s, opt_offset) {
                            (Some(s), Some(offset)) => {
                                Ok(Some(s.as_ref().gather_every(n as usize, offset as usize)?))
                            },
                            _ => Ok(None),
                        }
                    })?
                } else {
                    ListChunked::full_null_with_dtype(
                        list_ca.name().clone(),
                        list_ca.len(),
                        list_ca.inner_dtype(),
                    )
                }
            },
            (len_n, 1) if len_n == list_ca.len() => {
                if let Some(offset) = offset.get(0) {
                    list_ca.try_zip_and_apply_amortized(n, |opt_s, opt_n| match (opt_s, opt_n) {
                        (Some(s), Some(n)) => {
                            Ok(Some(s.as_ref().gather_every(n as usize, offset as usize)?))
                        },
                        _ => Ok(None),
                    })?
                } else {
                    ListChunked::full_null_with_dtype(
                        list_ca.name().clone(),
                        list_ca.len(),
                        list_ca.inner_dtype(),
                    )
                }
            },
            (len_n, len_offset) if len_n == len_offset && len_n == list_ca.len() => list_ca
                .try_binary_zip_and_apply_amortized(
                    n,
                    offset,
                    |opt_s, opt_n, opt_offset| match (opt_s, opt_n, opt_offset) {
                        (Some(s), Some(n), Some(offset)) => {
                            Ok(Some(s.as_ref().gather_every(n as usize, offset as usize)?))
                        },
                        _ => Ok(None),
                    },
                )?,
            _ => {
                polars_bail!(ComputeError: "The lengths of `n` and `offset` should be 1 or equal to the length of list.")
            },
        };
        Ok(out.into_series())
    }

    #[cfg(feature = "list_gather")]
    fn lst_gather(&self, idx: &Series, null_on_oob: bool) -> PolarsResult<Series> {
        let list_ca = self.as_list();
        let idx_ca = idx.list()?;

        polars_ensure!(
            idx_ca.inner_dtype().is_integer(),
            ComputeError: "cannot use dtype `{}` as an index", idx_ca.inner_dtype()
        );

        let index_typed_index = |idx: &Series| {
            let idx = idx.cast(&IDX_DTYPE).unwrap();
            {
                list_ca
                    .amortized_iter()
                    .map(|s| {
                        s.map(|s| {
                            let s = s.as_ref();
                            take_series(s, idx.clone(), null_on_oob)
                        })
                        .transpose()
                    })
                    .collect::<PolarsResult<ListChunked>>()
                    .map(|mut ca| {
                        ca.rename(list_ca.name().clone());
                        ca.into_series()
                    })
            }
        };

        match (list_ca.len(), idx_ca.len()) {
            (1, _) => {
                let mut out = if list_ca.has_nulls() {
                    ListChunked::full_null_with_dtype(
                        PlSmallStr::EMPTY,
                        idx.len(),
                        list_ca.inner_dtype(),
                    )
                } else {
                    let s = list_ca.explode(ExplodeOptions {
                        empty_as_null: true,
                        keep_nulls: true,
                    })?;
                    idx_ca
                        .series_iter()
                        .map(|opt_idx| {
                            opt_idx
                                .map(|idx| take_series(&s, idx, null_on_oob))
                                .transpose()
                        })
                        .collect::<PolarsResult<ListChunked>>()?
                };
                out.rename(list_ca.name().clone());
                Ok(out.into_series())
            },
            (_, 1) => {
                let idx_ca = idx_ca.explode(ExplodeOptions {
                    empty_as_null: true,
                    keep_nulls: true,
                })?;

                use DataType as D;
                match idx_ca.dtype() {
                    D::UInt32 | D::UInt64 => index_typed_index(&idx_ca),
                    dt if dt.is_signed_integer() => {
                        if let Some(min) = idx_ca.min::<i64>().unwrap() {
                            if min >= 0 {
                                index_typed_index(&idx_ca)
                            } else {
                                let mut out = {
                                    list_ca
                                        .amortized_iter()
                                        .map(|opt_s| {
                                            opt_s
                                                .map(|s| {
                                                    take_series(
                                                        s.as_ref(),
                                                        idx_ca.clone(),
                                                        null_on_oob,
                                                    )
                                                })
                                                .transpose()
                                        })
                                        .collect::<PolarsResult<ListChunked>>()?
                                };
                                out.rename(list_ca.name().clone());
                                Ok(out.into_series())
                            }
                        } else {
                            polars_bail!(ComputeError: "all indices are null");
                        }
                    },
                    dt => polars_bail!(ComputeError: "cannot use dtype `{dt}` as an index"),
                }
            },
            (a, b) if a == b => {
                let mut out = {
                    list_ca
                        .amortized_iter()
                        .zip(idx_ca.series_iter())
                        .map(|(opt_s, opt_idx)| {
                            {
                                match (opt_s, opt_idx) {
                                    (Some(s), Some(idx)) => {
                                        Some(take_series(s.as_ref(), idx, null_on_oob))
                                    },
                                    _ => None,
                                }
                            }
                            .transpose()
                        })
                        .collect::<PolarsResult<ListChunked>>()?
                };
                out.rename(list_ca.name().clone());
                Ok(out.into_series())
            },
            (a, b) => polars_bail!(length_mismatch = "list.gather", a, b),
        }
    }

    #[cfg(feature = "list_drop_nulls")]
    fn lst_drop_nulls(&self) -> ListChunked {
        let list_ca = self.as_list();

        // SAFETY: `drop_nulls` doesn't change the dtype
        unsafe { list_ca.apply_amortized_same_type(|s| s.as_ref().drop_nulls()) }
    }

    #[cfg(feature = "list_sample")]
    fn lst_sample_n(
        &self,
        n: &Series,
        with_replacement: bool,
        shuffle: Option<bool>,
        seed: Option<u64>,
    ) -> PolarsResult<ListChunked> {
        let ca = self.as_list();

        let n_s = n.strict_cast(&IDX_DTYPE)?;
        let n = n_s.idx()?;

        polars_ensure!(
            ca.len() == n.len() || ca.len() == 1 || n.len() == 1,
            length_mismatch = "list.sample(n)",
            ca.len(),
            n.len()
        );

        let target_len = n.len();
        if ca.len() == 1 && target_len > 1 {
            let single_list = ca.get_as_series(0);
            let out = sample_n_broadcast_list(
                single_list,
                n,
                with_replacement,
                shuffle,
                seed,
                target_len,
                ca.name().clone(),
                ca.inner_dtype(),
            )?;
            return Ok(self.same_type(out));
        }

        let out = match n.len() {
            1 => {
                if let Some(n) = n.get(0) {
                    unsafe {
                        // SAFETY: `sample_n` doesn't change the dtype
                        ca.try_apply_amortized_same_type(|s| {
                            s.as_ref()
                                .sample_n(n as usize, with_replacement, shuffle, seed)
                        })
                    }
                } else {
                    Ok(ListChunked::full_null_with_dtype(
                        ca.name().clone(),
                        ca.len(),
                        ca.inner_dtype(),
                    ))
                }
            },
            _ => ca.try_zip_and_apply_amortized(n, |opt_s, opt_n| match (opt_s, opt_n) {
                (Some(s), Some(n)) => s
                    .as_ref()
                    .sample_n(n as usize, with_replacement, shuffle, seed)
                    .map(Some),
                _ => Ok(None),
            }),
        };
        out.map(|ok| self.same_type(ok))
    }

    #[cfg(feature = "list_sample")]
    fn lst_sample_fraction(
        &self,
        fraction: &Series,
        with_replacement: bool,
        shuffle: Option<bool>,
        seed: Option<u64>,
    ) -> PolarsResult<ListChunked> {
        let ca = self.as_list();

        let fraction_s = fraction.cast(&DataType::Float64)?;
        let fraction = fraction_s.f64()?;

        if !with_replacement {
            for frac in fraction.iter().flatten() {
                polars_ensure!(
                    (0.0..=1.0).contains(&frac),
                    ComputeError: "fraction must be between 0.0 and 1.0, got: {}", frac
                )
            }
        }

        polars_ensure!(
            ca.len() == fraction.len() || ca.len() == 1 || fraction.len() == 1,
            length_mismatch = "list.sample(fraction)",
            ca.len(),
            fraction.len()
        );

        let target_len = fraction.len();
        if ca.len() == 1 && target_len > 1 {
            let single_list = ca.get_as_series(0);
            let out = sample_frac_broadcast_list(
                single_list,
                fraction,
                with_replacement,
                shuffle,
                seed,
                target_len,
                ca.name().clone(),
                ca.inner_dtype(),
            )?;
            return Ok(self.same_type(out));
        }

        let out = match fraction.len() {
            1 => {
                if let Some(fraction) = fraction.get(0) {
                    unsafe {
                        // SAFETY: `sample_n` doesn't change the dtype
                        ca.try_apply_amortized_same_type(|s| {
                            let n = (s.as_ref().len() as f64 * fraction) as usize;
                            s.as_ref().sample_n(n, with_replacement, shuffle, seed)
                        })
                    }
                } else {
                    Ok(ListChunked::full_null_with_dtype(
                        ca.name().clone(),
                        ca.len(),
                        ca.inner_dtype(),
                    ))
                }
            },
            _ => ca.try_zip_and_apply_amortized(fraction, |opt_s, opt_n| match (opt_s, opt_n) {
                (Some(s), Some(fraction)) => {
                    let n = (s.as_ref().len() as f64 * fraction) as usize;
                    s.as_ref()
                        .sample_n(n, with_replacement, shuffle, seed)
                        .map(Some)
                },
                _ => Ok(None),
            }),
        };
        out.map(|ok| self.same_type(ok))
    }

    fn lst_concat(&self, other: &[Column]) -> PolarsResult<ListChunked> {
        let ca = self.as_list();
        let other_len = other.len();
        let length = ca.len();
        let mut other = other.to_vec();
        let mut inner_super_type = ca.inner_dtype().clone();

        for s in &other {
            match s.dtype() {
                DataType::List(inner_type) => {
                    inner_super_type = try_get_supertype(&inner_super_type, inner_type)?;
                },
                dt => {
                    inner_super_type = try_get_supertype(&inner_super_type, dt)?;
                },
            }
        }

        // cast lhs
        let dtype = &DataType::List(Box::new(inner_super_type.clone()));
        let ca = ca.cast(dtype)?;
        let ca = ca.list().unwrap();

        // broadcasting path in case all unit length
        // this path will not expand the series, so saves memory
        let out = if other.iter().all(|s| s.len() == 1) && ca.len() != 1 {
            cast_rhs(&mut other, &inner_super_type, dtype, length, false)?;
            let to_append = other
                .iter()
                .filter_map(|s| {
                    let lst = s.list().unwrap();
                    // SAFETY: previous rhs_cast ensures the type is correct
                    unsafe {
                        lst.get_as_series(0)
                            .map(|s| s.from_physical_unchecked(&inner_super_type).unwrap())
                    }
                })
                .collect::<Vec<_>>();

            // there was a None, so all values will be None
            if to_append.len() != other_len {
                return Ok(ListChunked::full_null_with_dtype(
                    ca.name().clone(),
                    length,
                    &inner_super_type,
                ));
            }

            let vals_size_other = other
                .iter()
                .map(|s| s.list().unwrap().get_values_size())
                .sum::<usize>();

            let mut builder = get_list_builder(
                &inner_super_type,
                ca.get_values_size() + vals_size_other + 1,
                length,
                ca.name().clone(),
            );
            ca.series_iter().for_each(|opt_s| {
                let opt_s = opt_s.map(|mut s| {
                    for append in &to_append {
                        s.append(append).unwrap();
                    }
                    match inner_super_type {
                        // structs don't have chunks, so we must first rechunk the underlying series
                        #[cfg(feature = "dtype-struct")]
                        DataType::Struct(_) => s = s.rechunk(),
                        // nothing
                        _ => {},
                    }
                    s
                });
                builder.append_opt_series(opt_s.as_ref()).unwrap();
            });
            builder.finish()
        } else {
            // normal path which may contain same length list or unit length lists
            cast_rhs(&mut other, &inner_super_type, dtype, length, true)?;

            let vals_size_other = other
                .iter()
                .map(|s| s.list().unwrap().get_values_size())
                .sum::<usize>();
            let mut iters = Vec::with_capacity(other_len + 1);

            for s in other.iter_mut() {
                iters.push(s.list()?.amortized_iter())
            }
            let mut first_iter = ca.series_iter();
            let mut builder = get_list_builder(
                &inner_super_type,
                ca.get_values_size() + vals_size_other + 1,
                length,
                ca.name().clone(),
            );

            for _ in 0..ca.len() {
                let mut acc = match first_iter.next().unwrap() {
                    Some(s) => s,
                    None => {
                        builder.append_null();
                        // make sure that the iterators advance before we continue
                        for it in &mut iters {
                            it.next().unwrap();
                        }
                        continue;
                    },
                };

                let mut has_nulls = false;
                for it in &mut iters {
                    match it.next().unwrap() {
                        Some(s) => {
                            if !has_nulls {
                                acc.append(s.as_ref())?;
                            }
                        },
                        None => {
                            has_nulls = true;
                        },
                    }
                }
                if has_nulls {
                    builder.append_null();
                    continue;
                }

                match inner_super_type {
                    // structs don't have chunks, so we must first rechunk the underlying series
                    #[cfg(feature = "dtype-struct")]
                    DataType::Struct(_) => acc = acc.rechunk(),
                    // nothing
                    _ => {},
                }
                builder.append_series(&acc).unwrap();
            }
            builder.finish()
        };
        Ok(out)
    }
}

impl ListNameSpaceImpl for ListChunked {}

#[cfg(feature = "list_gather")]
fn take_series(s: &Series, idx: Series, null_on_oob: bool) -> PolarsResult<Series> {
    let len = s.len();
    let idx = convert_and_bound_index(&idx, len, null_on_oob)?;
    s.take(&idx)
}

pub fn slice_broadcast_list(
    single_list: Option<Series>,
    offsets: &Int64Chunked,
    lengths: &Int64Chunked,
    target_len: usize,
    name: PlSmallStr,
    inner_dtype: &DataType,
) -> ListChunked {
    debug_assert!(target_len == offsets.len().max(lengths.len()));

    let Some(single_list) = single_list else {
        return ListChunked::full_null_with_dtype(name, target_len, inner_dtype);
    };

    let iter = (0..target_len).map(|index| {
        let opt_offset = offsets.get(if offsets.len() == 1 { 0 } else { index });
        let opt_length = lengths.get(if lengths.len() == 1 { 0 } else { index });
        match (opt_offset, opt_length) {
            (Some(offset), Some(length)) => Some(single_list.slice(offset, length as usize)),
            _ => None,
        }
    });

    let mut out: ListChunked = iter.collect_trusted();
    out.rename(name);
    out
}

fn shift_broadcast_list(
    single_list: Option<Series>,
    periods: &Int64Chunked,
    target_len: usize,
    name: PlSmallStr,
    inner_dtype: &DataType,
) -> ListChunked {
    debug_assert!(target_len == periods.len());

    let Some(single_list) = single_list else {
        return ListChunked::full_null_with_dtype(name, target_len, inner_dtype);
    };

    let iter = (0..target_len).map(|index| {
        let opt_period = periods.get(index);
        opt_period.map(|period| single_list.shift(period))
    });

    let mut out: ListChunked = iter.collect_trusted();
    out.rename(name);
    out
}

#[cfg(feature = "list_sample")]
#[allow(clippy::too_many_arguments)]
fn sample_n_broadcast_list(
    single_list: Option<Series>,
    n: &IdxCa,
    with_replacement: bool,
    shuffle: Option<bool>,
    seed: Option<u64>,
    target_len: usize,
    name: PlSmallStr,
    inner_dtype: &DataType,
) -> PolarsResult<ListChunked> {
    debug_assert!(target_len == n.len());

    let Some(single_list) = single_list else {
        return Ok(ListChunked::full_null_with_dtype(
            name,
            target_len,
            inner_dtype,
        ));
    };

    let mut out: ListChunked = (0..target_len)
        .map(|index| -> PolarsResult<Option<Series>> {
            match n.get(index) {
                Some(n_val) => single_list
                    .sample_n(n_val as usize, with_replacement, shuffle, seed)
                    .map(Some),
                None => Ok(None),
            }
        })
        .collect::<PolarsResult<_>>()?;

    out.rename(name);
    Ok(out)
}

#[cfg(feature = "list_sample")]
#[allow(clippy::too_many_arguments)]
fn sample_frac_broadcast_list(
    single_list: Option<Series>,
    fraction: &Float64Chunked,
    with_replacement: bool,
    shuffle: Option<bool>,
    seed: Option<u64>,
    target_len: usize,
    name: PlSmallStr,
    inner_dtype: &DataType,
) -> PolarsResult<ListChunked> {
    debug_assert!(target_len == fraction.len());

    let Some(single_list) = single_list else {
        return Ok(ListChunked::full_null_with_dtype(
            name,
            target_len,
            inner_dtype,
        ));
    };

    let mut out: ListChunked = (0..target_len)
        .map(|index| -> PolarsResult<Option<Series>> {
            match fraction.get(index) {
                Some(frac_val) => single_list
                    .sample_frac(frac_val, with_replacement, shuffle, seed)
                    .map(Some),
                None => Ok(None),
            }
        })
        .collect::<PolarsResult<_>>()?;

    out.rename(name);
    Ok(out)
}