use polars_core::prelude::*;
#[cfg(feature = "array_to_struct")]
use polars_ops::chunked_array::array::{
arr_default_struct_name_gen, ArrToStructNameGenerator, ToStruct,
};
use crate::dsl::function_expr::ArrayFunction;
use crate::prelude::*;
pub struct ArrayNameSpace(pub Expr);
impl ArrayNameSpace {
pub fn max(self) -> Expr {
self.0
.map_private(FunctionExpr::ArrayExpr(ArrayFunction::Max))
}
pub fn min(self) -> Expr {
self.0
.map_private(FunctionExpr::ArrayExpr(ArrayFunction::Min))
}
pub fn sum(self) -> Expr {
self.0
.map_private(FunctionExpr::ArrayExpr(ArrayFunction::Sum))
}
pub fn std(self, ddof: u8) -> Expr {
self.0
.map_private(FunctionExpr::ArrayExpr(ArrayFunction::Std(ddof)))
}
pub fn var(self, ddof: u8) -> Expr {
self.0
.map_private(FunctionExpr::ArrayExpr(ArrayFunction::Var(ddof)))
}
pub fn median(self) -> Expr {
self.0
.map_private(FunctionExpr::ArrayExpr(ArrayFunction::Median))
}
pub fn unique(self) -> Expr {
self.0
.map_private(FunctionExpr::ArrayExpr(ArrayFunction::Unique(false)))
}
pub fn unique_stable(self) -> Expr {
self.0
.map_private(FunctionExpr::ArrayExpr(ArrayFunction::Unique(true)))
}
pub fn n_unique(self) -> Expr {
self.0
.map_private(FunctionExpr::ArrayExpr(ArrayFunction::NUnique))
}
pub fn to_list(self) -> Expr {
self.0
.map_private(FunctionExpr::ArrayExpr(ArrayFunction::ToList))
}
#[cfg(feature = "array_any_all")]
pub fn all(self) -> Expr {
self.0
.map_private(FunctionExpr::ArrayExpr(ArrayFunction::All))
}
#[cfg(feature = "array_any_all")]
pub fn any(self) -> Expr {
self.0
.map_private(FunctionExpr::ArrayExpr(ArrayFunction::Any))
}
pub fn sort(self, options: SortOptions) -> Expr {
self.0
.map_private(FunctionExpr::ArrayExpr(ArrayFunction::Sort(options)))
}
pub fn reverse(self) -> Expr {
self.0
.map_private(FunctionExpr::ArrayExpr(ArrayFunction::Reverse))
}
pub fn arg_min(self) -> Expr {
self.0
.map_private(FunctionExpr::ArrayExpr(ArrayFunction::ArgMin))
}
pub fn arg_max(self) -> Expr {
self.0
.map_private(FunctionExpr::ArrayExpr(ArrayFunction::ArgMax))
}
pub fn get(self, index: Expr, null_on_oob: bool) -> Expr {
self.0.map_many_private(
FunctionExpr::ArrayExpr(ArrayFunction::Get(null_on_oob)),
&[index],
false,
None,
)
}
pub fn join(self, separator: Expr, ignore_nulls: bool) -> Expr {
self.0.map_many_private(
FunctionExpr::ArrayExpr(ArrayFunction::Join(ignore_nulls)),
&[separator],
false,
None,
)
}
#[cfg(feature = "is_in")]
pub fn contains<E: Into<Expr>>(self, other: E) -> Expr {
let other = other.into();
self.0.map_many_private(
FunctionExpr::ArrayExpr(ArrayFunction::Contains),
&[other],
false,
None,
)
}
#[cfg(feature = "array_count")]
pub fn count_matches<E: Into<Expr>>(self, element: E) -> Expr {
let other = element.into();
self.0
.map_many_private(
FunctionExpr::ArrayExpr(ArrayFunction::CountMatches),
&[other],
false,
None,
)
.with_function_options(|mut options| {
options.flags |= FunctionFlags::INPUT_WILDCARD_EXPANSION;
options
})
}
#[cfg(feature = "array_to_struct")]
pub fn to_struct(self, name_generator: Option<ArrToStructNameGenerator>) -> PolarsResult<Expr> {
Ok(self
.0
.map(
move |s| {
s.array()?
.to_struct(name_generator.clone())
.map(|s| Some(s.into_column()))
},
GetOutput::map_dtype(move |dt: &DataType| {
let DataType::Array(inner, width) = dt else {
polars_bail!(InvalidOperation: "expected Array type, got: {}", dt)
};
let fields = (0..*width)
.map(|i| {
let name = arr_default_struct_name_gen(i);
Field::new(name, inner.as_ref().clone())
})
.collect();
Ok(DataType::Struct(fields))
}),
)
.with_fmt("arr.to_struct"))
}
pub fn shift(self, n: Expr) -> Expr {
self.0.map_many_private(
FunctionExpr::ArrayExpr(ArrayFunction::Shift),
&[n],
false,
None,
)
}
}