polars-ops 0.55.2

use num_traits::Bounded;
#[cfg(feature = "dtype-struct")]
use polars_core::chunked_array::ops::row_encode::_get_rows_encoded_ca;
use polars_core::prelude::arity::unary_elementwise_values;
use polars_core::prelude::*;
use polars_core::series::IsSorted;
use polars_core::with_match_physical_numeric_polars_type;
#[cfg(feature = "hash")]
use polars_utils::aliases::PlSeedableRandomStateQuality;
use polars_utils::total_ord::TotalOrd;

use crate::series::ops::SeriesSealed;

pub trait SeriesMethods: SeriesSealed {
    /// Create a [`DataFrame`] with the unique `values` of this [`Series`] and a column `"counts"`
    /// with dtype [`IdxType`]
    fn value_counts(
        &self,
        sort: bool,
        parallel: bool,
        name: PlSmallStr,
        normalize: bool,
    ) -> PolarsResult<DataFrame> {
        let s = self.as_series();
        polars_ensure!(
            s.name() != &name,
            Duplicate: "using `value_counts` on a column/series named '{}' would lead to duplicate \
            column names; change `name` to fix", name,
        );
        let groups = s.group_tuples(parallel, sort)?;
        let values = unsafe { s.agg_first(&groups) }
            .with_name(s.name().clone())
            .into();
        let counts = groups.group_count().with_name(name.clone());

        let counts = if normalize {
            let len = s.len() as f64;
            let counts: Float64Chunked =
                unary_elementwise_values(&counts, |count| count as f64 / len);
            counts.into_column()
        } else {
            counts.into_column()
        };

        let height = counts.len();
        let cols = vec![values, counts];
        let df = unsafe { DataFrame::new_unchecked(height, cols) };
        if sort {
            df.sort(
                [name],
                SortMultipleOptions::default()
                    .with_order_descending(true)
                    .with_multithreaded(parallel),
            )
        } else {
            Ok(df)
        }
    }

    #[cfg(feature = "hash")]
    fn hash(&self, build_hasher: PlSeedableRandomStateQuality) -> UInt64Chunked {
        let s = self.as_series();
        let mut h = vec![];
        s.0.vec_hash(build_hasher, &mut h).unwrap();
        UInt64Chunked::from_vec(s.name().clone(), h)
    }

    fn ensure_sorted_arg(&self, operation: &str) -> PolarsResult<()> {
        polars_ensure!(
            self.is_sorted(SortOptions::default())?,
            InvalidOperation: "argument in operation '{}' is not sorted, please sort the 'expr/series/column' first",
            operation
        );
        Ok(())
    }

    /// Checks if a [`Series`] is sorted with concrete options. Tries to fail fast.
    ///
    /// For inference of `descending` / `nulls_last`, see [`Self::is_sorted_any`].
    fn is_sorted(&self, options: SortOptions) -> PolarsResult<bool> {
        is_sorted_impl(self.as_series(), options)
    }

    fn is_sorted_any(
        &self,
        descending: Option<bool>,
        nulls_last: Option<bool>,
    ) -> PolarsResult<bool> {
        let s = self.as_series();
        let (descending, nulls_last) = resolve_sort_options(s, descending, nulls_last)?;
        // When an option could not be inferred the series is trivially sorted along that axis
        // (e.g. all non-null values equal, or no nulls), so any value works; default to `false`.
        let options = SortOptions {
            descending: descending.unwrap_or(false),
            nulls_last: nulls_last.unwrap_or(false),
            ..Default::default()
        };
        is_sorted_impl(s, options)
    }
}

fn is_sorted_impl(s: &Series, options: SortOptions) -> PolarsResult<bool> {
    let null_count = s.null_count();

    if (options.descending
        && (options.nulls_last || null_count == 0)
        && matches!(s.is_sorted_flag(), IsSorted::Descending))
        || (!options.descending
            && (!options.nulls_last || null_count == 0)
            && matches!(s.is_sorted_flag(), IsSorted::Ascending))
    {
        return Ok(true);
    }

    #[cfg(feature = "dtype-struct")]
    if matches!(s.dtype(), DataType::Struct(_)) {
        let encoded = _get_rows_encoded_ca(
            PlSmallStr::EMPTY,
            &[s.clone().into()],
            &[options.descending],
            &[options.nulls_last],
            false,
        )?;
        let options = SortOptions {
            descending: false,
            nulls_last: false,
            ..options
        };
        return is_sorted_impl(&encoded.into_series(), options);
    }

    let s_len = s.len();
    if null_count == s_len {
        // All nulls are equal.
        return Ok(true);
    }
    // Check if nulls are in the right location.
    if null_count > 0 {
        if options.nulls_last {
            if s.slice((s_len - null_count) as i64, null_count)
                .null_count()
                != null_count
            {
                return Ok(false);
            }
        } else if s.slice(0, null_count).null_count() != null_count {
            return Ok(false);
        }
    }

    if s.dtype().is_primitive_numeric() {
        with_match_physical_numeric_polars_type!(s.dtype(), |$T| {
            let ca: &ChunkedArray<$T> = s.as_ref().as_ref().as_ref();
            return Ok(is_sorted_ca_num::<$T>(ca, options))
        })
    }

    // Logical non-primitive types (e.g. String, Categorical, List, …): take only the contiguous
    // non-null values (`non_null`). For ordinary `Categorical` use `iter_str` (below); otherwise
    // `to_physical_repr`, then
    // (1) for ordinary [`DataType::Categorical`], compare adjacent **decoded strings** (`iter_str`),
    // (2) reuse `is_sorted_ca_num` when the physical type is primitive numeric (temporal /
    //     Decimal, Enum-as-integer, …) after `to_physical_repr`;
    // (3) uses a dedicated kernel for boolean values,
    // (4) else scans string / binary values with `TotalOrd`,
    // (5) else fall back to pairwise `Series::lt_eq` / `gt_eq` (nested types, etc.).
    let non_null_len = s_len - null_count;
    if non_null_len <= 1 {
        return Ok(true);
    }

    let offset = (!options.nulls_last as i64) * (null_count as i64);
    let non_null = s.slice(offset, non_null_len);
    debug_assert_eq!(
        non_null.null_count(),
        0,
        "internal error: `is_sorted` non-null slice contains nulls"
    );

    #[cfg(feature = "dtype-categorical")]
    if matches!(non_null.dtype(), DataType::Categorical(_, _)) {
        return is_sorted_categorical_lexical_adjacent(&non_null, options);
    }

    let phys = non_null.to_physical_repr();
    let s_phys = phys.as_ref();
    if s_phys.dtype().is_primitive_numeric() {
        with_match_physical_numeric_polars_type!(s_phys.dtype(), |$T| {
            let ca: &ChunkedArray<$T> = s_phys.as_ref().as_ref().as_ref();
            return Ok(is_sorted_ca_num::<$T>(ca, options))
        })
    }

    match s_phys.dtype() {
        DataType::Boolean => {
            let ca = s_phys.bool()?;
            Ok(is_sorted_ca_bool(ca, options.descending))
        },
        DataType::String => {
            let ca = s_phys.str()?;
            Ok(is_sorted_adjacent_total_ord(
                ca.no_null_iter(),
                options.descending,
            ))
        },
        DataType::Binary => {
            let ca = s_phys.binary()?;
            Ok(is_sorted_adjacent_total_ord(
                ca.no_null_iter(),
                options.descending,
            ))
        },
        DataType::BinaryOffset => {
            let ca = s_phys.binary_offset()?;
            Ok(is_sorted_adjacent_total_ord(
                ca.no_null_iter(),
                options.descending,
            ))
        },
        _ => {
            // `non_null` excludes nulls already; compare `non_null[..-1]` with `non_null[1..]`.
            let cmp_len = non_null_len - 1;
            let s1 = non_null.slice(0, cmp_len);
            let s2 = non_null.slice(1, cmp_len);
            let cmp_op = if options.descending {
                Series::gt_eq
            } else {
                Series::lt_eq
            };
            Ok(cmp_op(&s1, &s2)?.all())
        },
    }
}

/// Returns whether iterator elements are non-decreasing (`descending == false`) or non-increasing
/// (`descending == true`) under [`TotalOrd`].
///
/// Assumes the iterator `it` yields **only** the non-null values in row order (one item per row). An empty
/// iterator is considered sorted. Stops at the first pair that violates the ordering.
fn is_sorted_adjacent_total_ord<T: TotalOrd>(
    it: impl Iterator<Item = T>,
    descending: bool,
) -> bool {
    let mut it = it;
    // Sliding window: `prev` is always the previous element; seed with the first value.
    let Some(mut prev) = it.next() else {
        return true;
    };
    if descending {
        for v in it {
            if !prev.tot_ge(&v) {
                return false;
            }
            prev = v;
        }
    } else {
        for v in it {
            if !prev.tot_le(&v) {
                return false;
            }
            prev = v;
        }
    }
    true
}

/// Ordinary [`DataType::Categorical`]: lexical order via adjacent decoded strings (`iter_str`), same as
/// `Series::lt_eq` / `gt_eq`, but without a Boolean series. Caller must pass a contiguous **non-null**
/// slice.
#[cfg(feature = "dtype-categorical")]
fn is_sorted_categorical_lexical_adjacent(s: &Series, options: SortOptions) -> PolarsResult<bool> {
    polars_ensure!(
        matches!(s.dtype(), DataType::Categorical(_, _)),
        ComputeError: "internal error: expected Categorical in lexical `is_sorted` path",
    );

    with_match_categorical_physical_type!(s.dtype().cat_physical().unwrap(), |$C| {
        let ca = s.cat::<$C>()?;

        // `ca.null_count() == 0` implies each `phys` row decodes via `iter_str` to `Some(..)`
        Ok(is_sorted_adjacent_total_ord(
            ca.iter_str().map(|opt| {
                opt.expect(
                    "`iter_str` produced None while categorical null_count reported 0 (`is_sorted`)"
                )
            }),
            options.descending,
        ))
    })
}

/// Booleans ordered as [`false`] < [`true`] (same as inequality comparisons on [`BooleanChunked`]).
///
/// Monotone order is equivalent to at most one plateau change: ascending is `F…FT…T`, descending is
/// `T…TF…F`. Implemented with `first_true_idx` / `first_false_idx` plus a global false/true count
/// check.
///
/// Caller must ensure **`ca` has no nulls** on the flattened series (see `non_null` slice above).
fn is_sorted_ca_bool(ca: &BooleanChunked, descending: bool) -> bool {
    let len = ca.len();
    if len <= 1 {
        return true;
    }
    debug_assert_eq!(
        ca.null_count(),
        0,
        "internal error: `is_sorted_ca_bool` expects a non-null boolean slice"
    );
    if descending {
        let Some(idx) = ca.first_false_idx() else {
            return true;
        };
        !ca.slice(idx as i64, ca.len() - idx).any()
    } else {
        let Some(idx) = ca.first_true_idx() else {
            return true;
        };
        ca.slice(idx as i64, ca.len() - idx).all()
    }
}

/// Infers the `(descending, nulls_last)` sort options for `s`, honoring any provided hints.
///
/// Each returned value is `Some` when known — taken from the corresponding hint when given,
/// otherwise inferred from the data — and `None` when it cannot be inferred from `s` alone:
/// - `descending` is `None` when there are fewer than two distinct non-null values, so no direction
///   is implied.
/// - `nulls_last` is `None` when `s` has no nulls, is entirely null, or the nulls are interleaved
///   (the last of which is not sorted under any placement and is rejected by the `is_sorted` check).
///
/// The two axes are independent, so callers can use whichever was determined even when the other
/// could not be.
pub fn resolve_sort_options(
    s: &Series,
    descending: Option<bool>,
    nulls_last: Option<bool>,
) -> PolarsResult<(Option<bool>, Option<bool>)> {
    let nulls_last = match nulls_last {
        Some(n) => Some(n),
        None => infer_nulls_last(s),
    };

    let descending = match descending {
        Some(d) => Some(d),
        None => infer_descending(s, nulls_last.unwrap_or(false))?,
    };

    Ok((descending, nulls_last))
}

/// Infers null placement from `s`: `Some(true)` if all nulls sit at the tail, `Some(false)` if all
/// sit at the head, and `None` if there are no nulls, `s` is entirely null, or the nulls are
/// interleaved (the latter is not sorted under any placement; the `is_sorted` check rejects it).
fn infer_nulls_last(s: &Series) -> Option<bool> {
    let null_count = s.null_count();
    let s_len = s.len();

    if null_count == 0 || null_count == s_len {
        return None;
    }

    if s.slice((s_len - null_count) as i64, null_count)
        .null_count()
        == null_count
    {
        Some(true)
    } else if s.slice(0, null_count).null_count() == null_count {
        Some(false)
    } else {
        None
    }
}

fn infer_descending(s: &Series, nulls_last: bool) -> PolarsResult<Option<bool>> {
    let null_count = s.null_count();
    let non_null_len = s.len() - null_count;
    if non_null_len < 2 {
        return Ok(None);
    }

    let non_null_start = if nulls_last { 0 } else { null_count };
    let non_null = s.slice(non_null_start as i64, non_null_len);

    let a = non_null.slice(0, non_null_len - 1);
    let b = non_null.slice(1, non_null_len - 1);

    let lt = a.lt(&b)?;
    let gt = a.gt(&b)?;

    let lt_first = lt.iter().position(|v| v == Some(true));
    let gt_first = gt.iter().position(|v| v == Some(true));

    Ok(match (lt_first, gt_first) {
        (None, None) => None,
        (Some(_), None) => Some(false),
        (None, Some(_)) => Some(true),
        (Some(l), Some(g)) => Some(g < l),
    })
}

fn check_cmp<T: NumericNative, Cmp: Fn(&T, &T) -> bool>(
    vals: &[T],
    f: Cmp,
    previous: &mut T,
) -> bool {
    let mut sorted = true;
    for c in vals.chunks(1024) {
        for v in c {
            sorted &= f(previous, v);
            *previous = *v;
        }
        if !sorted {
            return false;
        }
    }
    sorted
}

fn is_sorted_ca_num<T: PolarsNumericType>(ca: &ChunkedArray<T>, options: SortOptions) -> bool {
    if let Ok(vals) = ca.cont_slice() {
        let mut previous = vals[0];
        return if options.descending {
            check_cmp(vals, |prev, c| prev.tot_ge(c), &mut previous)
        } else {
            check_cmp(vals, |prev, c| prev.tot_le(c), &mut previous)
        };
    };

    if ca.null_count() == 0 {
        let mut previous = if options.descending {
            T::Native::max_value()
        } else {
            T::Native::min_value()
        };
        for arr in ca.downcast_iter() {
            let vals = arr.values();
            let sorted = if options.descending {
                check_cmp(vals, |prev, c| prev.tot_ge(c), &mut previous)
            } else {
                check_cmp(vals, |prev, c| prev.tot_le(c), &mut previous)
            };
            if !sorted {
                return false;
            }
        }
        return true;
    };

    let null_count = ca.null_count();
    if options.nulls_last {
        let ca = ca.slice(0, ca.len() - null_count);
        is_sorted_ca_num(&ca, options)
    } else {
        let ca = ca.slice(null_count as i64, ca.len() - null_count);
        is_sorted_ca_num(&ca, options)
    }
}

impl SeriesMethods for Series {}