diff --git a/crates/polars-core/src/chunked_array/arithmetic/mod.rs b/crates/polars-core/src/chunked_array/arithmetic/mod.rs index c39d9edb26b1..ab4b3e7eb337 100644 --- a/crates/polars-core/src/chunked_array/arithmetic/mod.rs +++ b/crates/polars-core/src/chunked_array/arithmetic/mod.rs @@ -9,6 +9,7 @@ use arrow::compute::utils::combine_validities_and; use num_traits::{Num, NumCast, ToPrimitive}; pub use numeric::ArithmeticChunked; +use crate::prelude::arity::unary_elementwise_values; use crate::prelude::*; #[inline] @@ -135,7 +136,7 @@ impl Add for &BooleanChunked { if rhs.len() == 1 { let rhs = rhs.get(0); return match rhs { - Some(rhs) => self.apply_values_generic(|v| v as IdxSize + rhs as IdxSize), + Some(rhs) => unary_elementwise_values(self, |v| v as IdxSize + rhs as IdxSize), None => IdxCa::full_null(self.name(), self.len()), }; } diff --git a/crates/polars-core/src/chunked_array/float.rs b/crates/polars-core/src/chunked_array/float.rs index e2c7d726774d..dc09024b704c 100644 --- a/crates/polars-core/src/chunked_array/float.rs +++ b/crates/polars-core/src/chunked_array/float.rs @@ -3,6 +3,7 @@ use arrow::legacy::kernels::set::set_at_nulls; use num_traits::Float; use polars_utils::total_ord::{canonical_f32, canonical_f64}; +use crate::prelude::arity::unary_elementwise_values; use crate::prelude::*; impl ChunkedArray @@ -57,6 +58,6 @@ where T::Native: Float + Canonical, { pub fn to_canonical(&self) -> Self { - self.apply_values_generic(|v| v.canonical()) + unary_elementwise_values(self, |v| v.canonical()) } } diff --git a/crates/polars-core/src/chunked_array/mod.rs b/crates/polars-core/src/chunked_array/mod.rs index fe145f7c8b76..71625393ec51 100644 --- a/crates/polars-core/src/chunked_array/mod.rs +++ b/crates/polars-core/src/chunked_array/mod.rs @@ -75,8 +75,8 @@ pub type ChunkLenIter<'a> = std::iter::Map, fn(&A /// /// ```rust /// # use polars_core::prelude::*; -/// fn apply_cosine_and_cast(ca: &Float32Chunked) -> Float64Chunked { -/// ca.apply_values_generic(|v| v.cos() as f64) +/// fn apply_cosine_and_cast(ca: &Float32Chunked) -> Float32Chunked { +/// ca.apply_values(|v| v.cos()) /// } /// ``` /// diff --git a/crates/polars-core/src/chunked_array/ops/aggregate/var.rs b/crates/polars-core/src/chunked_array/ops/aggregate/var.rs index aff0aeb69641..1ca04cc2b30e 100644 --- a/crates/polars-core/src/chunked_array/ops/aggregate/var.rs +++ b/crates/polars-core/src/chunked_array/ops/aggregate/var.rs @@ -1,3 +1,5 @@ +use arity::unary_elementwise_values; + use super::*; pub trait VarAggSeries { @@ -19,7 +21,7 @@ where } let mean = self.mean()?; - let squared: Float64Chunked = ChunkedArray::apply_values_generic(self, |value| { + let squared: Float64Chunked = unary_elementwise_values(self, |value| { let tmp = value.to_f64().unwrap() - mean; tmp * tmp }); diff --git a/crates/polars-core/src/chunked_array/ops/apply.rs b/crates/polars-core/src/chunked_array/ops/apply.rs index 09e0572b89d7..7119a48f58d3 100644 --- a/crates/polars-core/src/chunked_array/ops/apply.rs +++ b/crates/polars-core/src/chunked_array/ops/apply.rs @@ -1,6 +1,7 @@ //! Implementations of the ChunkApply Trait. use std::borrow::Cow; +use crate::chunked_array::arity::{unary_elementwise, unary_elementwise_values}; use crate::chunked_array::cast::CastOptions; use crate::prelude::*; use crate::series::IsSorted; @@ -9,23 +10,6 @@ impl ChunkedArray where T: PolarsDataType, { - // Applies a function to all elements, regardless of whether they - // are null or not, after which the null mask is copied from the - // original array. - pub fn apply_values_generic<'a, U, K, F>(&'a self, mut op: F) -> ChunkedArray - where - U: PolarsDataType, - F: FnMut(T::Physical<'a>) -> K, - U::Array: ArrayFromIter, - { - let iter = self.downcast_iter().map(|arr| { - let out: U::Array = arr.values_iter().map(&mut op).collect_arr(); - out.with_validity_typed(arr.validity().cloned()) - }); - - ChunkedArray::from_chunk_iter(self.name(), iter) - } - /// Applies a function only to the non-null elements, propagating nulls. pub fn apply_nonnull_values_generic<'a, U, K, F>( &'a self, @@ -83,39 +67,6 @@ where ChunkedArray::try_from_chunk_iter(self.name(), iter) } - pub fn apply_generic<'a, U, K, F>(&'a self, mut op: F) -> ChunkedArray - where - U: PolarsDataType, - F: FnMut(Option>) -> Option, - U::Array: ArrayFromIter>, - { - if self.null_count() == 0 { - let iter = self - .downcast_iter() - .map(|arr| arr.values_iter().map(|x| op(Some(x))).collect_arr()); - ChunkedArray::from_chunk_iter(self.name(), iter) - } else { - let iter = self - .downcast_iter() - .map(|arr| arr.iter().map(&mut op).collect_arr()); - ChunkedArray::from_chunk_iter(self.name(), iter) - } - } - - pub fn try_apply_generic<'a, U, K, F, E>(&'a self, op: F) -> Result, E> - where - U: PolarsDataType, - F: FnMut(Option>) -> Result, E> + Copy, - U::Array: ArrayFromIter>, - { - let iter = self.downcast_iter().map(|arr| { - let array: U::Array = arr.iter().map(op).try_collect_arr()?; - Ok(array.with_validity_typed(arr.validity().cloned())) - }); - - ChunkedArray::try_from_chunk_iter(self.name(), iter) - } - pub fn apply_into_string_amortized<'a, F>(&'a self, mut f: F) -> StringChunked where F: FnMut(T::Physical<'a>, &mut String), @@ -329,7 +280,7 @@ impl<'a> ChunkApply<'a, bool> for BooleanChunked { where F: Fn(Option) -> Option + Copy, { - self.apply_generic(f) + unary_elementwise(self, f) } fn apply_to_slice(&'a self, f: F, slice: &mut [T]) @@ -386,14 +337,14 @@ impl<'a> ChunkApply<'a, &'a str> for StringChunked { where F: Fn(&'a str) -> Cow<'a, str> + Copy, { - ChunkedArray::apply_values_generic(self, f) + unary_elementwise_values(self, f) } fn apply(&'a self, f: F) -> Self where F: Fn(Option<&'a str>) -> Option> + Copy, { - self.apply_generic(f) + unary_elementwise(self, f) } fn apply_to_slice(&'a self, f: F, slice: &mut [T]) @@ -422,14 +373,14 @@ impl<'a> ChunkApply<'a, &'a [u8]> for BinaryChunked { where F: Fn(&'a [u8]) -> Cow<'a, [u8]> + Copy, { - self.apply_values_generic(f) + unary_elementwise_values(self, f) } fn apply(&'a self, f: F) -> Self where F: Fn(Option<&'a [u8]>) -> Option> + Copy, { - self.apply_generic(f) + unary_elementwise(self, f) } fn apply_to_slice(&'a self, f: F, slice: &mut [T]) diff --git a/crates/polars-core/src/chunked_array/ops/arity.rs b/crates/polars-core/src/chunked_array/ops/arity.rs index a28c5dcf6344..c69ae14a5866 100644 --- a/crates/polars-core/src/chunked_array/ops/arity.rs +++ b/crates/polars-core/src/chunked_array/ops/arity.rs @@ -74,10 +74,17 @@ where F: UnaryFnMut>>, V::Array: ArrayFromIter<>>>::Ret>, { - let iter = ca - .downcast_iter() - .map(|arr| arr.iter().map(&mut op).collect_arr()); - ChunkedArray::from_chunk_iter(ca.name(), iter) + if ca.has_nulls() { + let iter = ca + .downcast_iter() + .map(|arr| arr.iter().map(&mut op).collect_arr()); + ChunkedArray::from_chunk_iter(ca.name(), iter) + } else { + let iter = ca + .downcast_iter() + .map(|arr| arr.values_iter().map(|x| op(Some(x))).collect_arr()); + ChunkedArray::from_chunk_iter(ca.name(), iter) + } } #[inline] diff --git a/crates/polars-ops/src/chunked_array/binary/namespace.rs b/crates/polars-ops/src/chunked_array/binary/namespace.rs index c74ad116ab1e..31f92424d645 100644 --- a/crates/polars-ops/src/chunked_array/binary/namespace.rs +++ b/crates/polars-ops/src/chunked_array/binary/namespace.rs @@ -6,7 +6,7 @@ use base64::engine::general_purpose; #[cfg(feature = "binary_encoding")] use base64::Engine as _; use memchr::memmem::find; -use polars_core::prelude::arity::broadcast_binary_elementwise_values; +use polars_core::prelude::arity::{broadcast_binary_elementwise_values, unary_elementwise_values}; use super::*; @@ -15,7 +15,7 @@ pub trait BinaryNameSpaceImpl: AsBinary { fn contains(&self, lit: &[u8]) -> BooleanChunked { let ca = self.as_binary(); let f = |s: &[u8]| find(s, lit).is_some(); - ca.apply_values_generic(f) + unary_elementwise_values(ca, f) } fn contains_chunked(&self, lit: &BinaryChunked) -> BooleanChunked { diff --git a/crates/polars-ops/src/chunked_array/strings/concat.rs b/crates/polars-ops/src/chunked_array/strings/concat.rs index ea557267a1d2..67d1f244843d 100644 --- a/crates/polars-ops/src/chunked_array/strings/concat.rs +++ b/crates/polars-ops/src/chunked_array/strings/concat.rs @@ -1,5 +1,6 @@ use arrow::array::{Utf8Array, ValueSize}; use arrow::compute::cast::utf8_to_utf8view; +use polars_core::prelude::arity::unary_elementwise; use polars_core::prelude::*; // Vertically concatenate all strings in a StringChunked. @@ -67,7 +68,7 @@ pub fn hor_str_concat( return if !ignore_nulls || ca.null_count() == 0 { Ok(ca.clone()) } else { - Ok(ca.apply_generic(|val| Some(val.unwrap_or("")))) + Ok(unary_elementwise(ca, |val| Some(val.unwrap_or("")))) }; } diff --git a/crates/polars-ops/src/chunked_array/strings/find_many.rs b/crates/polars-ops/src/chunked_array/strings/find_many.rs index 8a1ef7815117..9bf0510e93d9 100644 --- a/crates/polars-ops/src/chunked_array/strings/find_many.rs +++ b/crates/polars-ops/src/chunked_array/strings/find_many.rs @@ -1,5 +1,6 @@ use aho_corasick::{AhoCorasick, AhoCorasickBuilder}; use arrow::array::Utf8ViewArray; +use polars_core::prelude::arity::unary_elementwise; use polars_core::prelude::*; use polars_core::utils::align_chunks_binary; @@ -27,7 +28,9 @@ pub fn contains_any( ) -> PolarsResult { let ac = build_ac(patterns, ascii_case_insensitive)?; - Ok(ca.apply_generic(|opt_val| opt_val.map(|val| ac.find(val).is_some()))) + Ok(unary_elementwise(ca, |opt_val| { + opt_val.map(|val| ac.find(val).is_some()) + })) } pub fn replace_all( @@ -52,7 +55,9 @@ pub fn replace_all( let ac = build_ac(patterns, ascii_case_insensitive)?; - Ok(ca.apply_generic(|opt_val| opt_val.map(|val| ac.replace_all(val, replace_with.as_slice())))) + Ok(unary_elementwise(ca, |opt_val| { + opt_val.map(|val| ac.replace_all(val, replace_with.as_slice())) + })) } fn push(val: &str, builder: &mut ListStringChunkedBuilder, ac: &AhoCorasick, overlapping: bool) { diff --git a/crates/polars-ops/src/chunked_array/strings/namespace.rs b/crates/polars-ops/src/chunked_array/strings/namespace.rs index 099a4953fa17..894b36d3c84b 100644 --- a/crates/polars-ops/src/chunked_array/strings/namespace.rs +++ b/crates/polars-ops/src/chunked_array/strings/namespace.rs @@ -265,7 +265,7 @@ pub trait StringNameSpaceImpl: AsString { let res_reg = Regex::new(pat); let opt_reg = if strict { Some(res_reg?) } else { res_reg.ok() }; let out: BooleanChunked = if let Some(reg) = opt_reg { - ca.apply_values_generic(|s| reg.is_match(s)) + unary_elementwise_values(ca, |s| reg.is_match(s)) } else { BooleanChunked::full_null(ca.name(), ca.len()) }; @@ -289,11 +289,9 @@ pub trait StringNameSpaceImpl: AsString { fn find(&self, pat: &str, strict: bool) -> PolarsResult { let ca = self.as_string(); match Regex::new(pat) { - Ok(rx) => { - Ok(ca.apply_generic(|opt_s| { - opt_s.and_then(|s| rx.find(s)).map(|m| m.start() as u32) - })) - }, + Ok(rx) => Ok(unary_elementwise(ca, |opt_s| { + opt_s.and_then(|s| rx.find(s)).map(|m| m.start() as u32) + })), Err(_) if !strict => Ok(UInt32Chunked::full_null(ca.name(), ca.len())), Err(e) => Err(PolarsError::ComputeError( format!("Invalid regular expression: {}", e).into(), @@ -419,7 +417,7 @@ pub trait StringNameSpaceImpl: AsString { fn strip_chars(&self, pat: &Series) -> PolarsResult { let ca = self.as_string(); if pat.dtype() == &DataType::Null { - Ok(ca.apply_generic(|opt_s| opt_s.map(|s| s.trim()))) + Ok(unary_elementwise(ca, |opt_s| opt_s.map(|s| s.trim()))) } else { Ok(strip_chars(ca, pat.str()?)) } @@ -428,7 +426,7 @@ pub trait StringNameSpaceImpl: AsString { fn strip_chars_start(&self, pat: &Series) -> PolarsResult { let ca = self.as_string(); if pat.dtype() == &DataType::Null { - return Ok(ca.apply_generic(|opt_s| opt_s.map(|s| s.trim_start()))); + return Ok(unary_elementwise(ca, |opt_s| opt_s.map(|s| s.trim_start()))); } else { Ok(strip_chars_start(ca, pat.str()?)) } @@ -437,7 +435,7 @@ pub trait StringNameSpaceImpl: AsString { fn strip_chars_end(&self, pat: &Series) -> PolarsResult { let ca = self.as_string(); if pat.dtype() == &DataType::Null { - return Ok(ca.apply_generic(|opt_s| opt_s.map(|s| s.trim_end()))); + return Ok(unary_elementwise(ca, |opt_s| opt_s.map(|s| s.trim_end()))); } else { Ok(strip_chars_end(ca, pat.str()?)) } @@ -524,7 +522,9 @@ pub trait StringNameSpaceImpl: AsString { Regex::new(pat)? }; - Ok(ca.apply_generic(|opt_s| opt_s.map(|s| reg.find_iter(s).count() as u32))) + Ok(unary_elementwise(ca, |opt_s| { + opt_s.map(|s| reg.find_iter(s).count() as u32) + })) } /// Count all successive non-overlapping regex matches. diff --git a/crates/polars-ops/src/chunked_array/strings/reverse.rs b/crates/polars-ops/src/chunked_array/strings/reverse.rs index cb32dbf245e7..763040d1cb2f 100644 --- a/crates/polars-ops/src/chunked_array/strings/reverse.rs +++ b/crates/polars-ops/src/chunked_array/strings/reverse.rs @@ -1,3 +1,4 @@ +use polars_core::prelude::arity::unary_elementwise; use polars_core::prelude::StringChunked; use unicode_reverse::reverse_grapheme_clusters_in_place; @@ -10,5 +11,5 @@ fn to_reverse_helper(s: Option<&str>) -> Option { } pub fn reverse(ca: &StringChunked) -> StringChunked { - ca.apply_generic(to_reverse_helper) + unary_elementwise(ca, to_reverse_helper) } diff --git a/crates/polars-ops/src/chunked_array/strings/strip.rs b/crates/polars-ops/src/chunked_array/strings/strip.rs index 5c671ee3bc7c..c7468d238807 100644 --- a/crates/polars-ops/src/chunked_array/strings/strip.rs +++ b/crates/polars-ops/src/chunked_array/strings/strip.rs @@ -1,4 +1,4 @@ -use polars_core::prelude::arity::broadcast_binary_elementwise; +use polars_core::prelude::arity::{broadcast_binary_elementwise, unary_elementwise}; use super::*; @@ -58,14 +58,16 @@ pub fn strip_chars(ca: &StringChunked, pat: &StringChunked) -> StringChunked { if let Some(pat) = pat.get(0) { if pat.chars().count() == 1 { // Fast path for when a single character is passed - ca.apply_generic(|opt_s| { + unary_elementwise(ca, |opt_s| { opt_s.map(|s| s.trim_matches(pat.chars().next().unwrap())) }) } else { - ca.apply_generic(|opt_s| opt_s.map(|s| s.trim_matches(|c| pat.contains(c)))) + unary_elementwise(ca, |opt_s| { + opt_s.map(|s| s.trim_matches(|c| pat.contains(c))) + }) } } else { - ca.apply_generic(|opt_s| opt_s.map(|s| s.trim())) + unary_elementwise(ca, |opt_s| opt_s.map(|s| s.trim())) } }, _ => broadcast_binary_elementwise(ca, pat, strip_chars_binary), @@ -78,16 +80,16 @@ pub fn strip_chars_start(ca: &StringChunked, pat: &StringChunked) -> StringChunk if let Some(pat) = pat.get(0) { if pat.chars().count() == 1 { // Fast path for when a single character is passed - ca.apply_generic(|opt_s| { + unary_elementwise(ca, |opt_s| { opt_s.map(|s| s.trim_start_matches(pat.chars().next().unwrap())) }) } else { - ca.apply_generic(|opt_s| { + unary_elementwise(ca, |opt_s| { opt_s.map(|s| s.trim_start_matches(|c| pat.contains(c))) }) } } else { - ca.apply_generic(|opt_s| opt_s.map(|s| s.trim_start())) + unary_elementwise(ca, |opt_s| opt_s.map(|s| s.trim_start())) } }, _ => broadcast_binary_elementwise(ca, pat, strip_chars_start_binary), @@ -100,14 +102,16 @@ pub fn strip_chars_end(ca: &StringChunked, pat: &StringChunked) -> StringChunked if let Some(pat) = pat.get(0) { if pat.chars().count() == 1 { // Fast path for when a single character is passed - ca.apply_generic(|opt_s| { + unary_elementwise(ca, |opt_s| { opt_s.map(|s| s.trim_end_matches(pat.chars().next().unwrap())) }) } else { - ca.apply_generic(|opt_s| opt_s.map(|s| s.trim_end_matches(|c| pat.contains(c)))) + unary_elementwise(ca, |opt_s| { + opt_s.map(|s| s.trim_end_matches(|c| pat.contains(c))) + }) } } else { - ca.apply_generic(|opt_s| opt_s.map(|s| s.trim_end())) + unary_elementwise(ca, |opt_s| opt_s.map(|s| s.trim_end())) } }, _ => broadcast_binary_elementwise(ca, pat, strip_chars_end_binary), @@ -117,9 +121,9 @@ pub fn strip_chars_end(ca: &StringChunked, pat: &StringChunked) -> StringChunked pub fn strip_prefix(ca: &StringChunked, prefix: &StringChunked) -> StringChunked { match prefix.len() { 1 => match prefix.get(0) { - Some(prefix) => { - ca.apply_generic(|opt_s| opt_s.map(|s| s.strip_prefix(prefix).unwrap_or(s))) - }, + Some(prefix) => unary_elementwise(ca, |opt_s| { + opt_s.map(|s| s.strip_prefix(prefix).unwrap_or(s)) + }), _ => StringChunked::full_null(ca.name(), ca.len()), }, _ => broadcast_binary_elementwise(ca, prefix, strip_prefix_binary), @@ -129,9 +133,9 @@ pub fn strip_prefix(ca: &StringChunked, prefix: &StringChunked) -> StringChunked pub fn strip_suffix(ca: &StringChunked, suffix: &StringChunked) -> StringChunked { match suffix.len() { 1 => match suffix.get(0) { - Some(suffix) => { - ca.apply_generic(|opt_s| opt_s.map(|s| s.strip_suffix(suffix).unwrap_or(s))) - }, + Some(suffix) => unary_elementwise(ca, |opt_s| { + opt_s.map(|s| s.strip_suffix(suffix).unwrap_or(s)) + }), _ => StringChunked::full_null(ca.name(), ca.len()), }, _ => broadcast_binary_elementwise(ca, suffix, strip_suffix_binary), diff --git a/crates/polars-ops/src/series/ops/clip.rs b/crates/polars-ops/src/series/ops/clip.rs index 37fc8b3030e8..485b2d7a064b 100644 --- a/crates/polars-ops/src/series/ops/clip.rs +++ b/crates/polars-ops/src/series/ops/clip.rs @@ -1,4 +1,4 @@ -use polars_core::prelude::arity::{binary_elementwise, ternary_elementwise}; +use polars_core::prelude::arity::{binary_elementwise, ternary_elementwise, unary_elementwise}; use polars_core::prelude::*; use polars_core::with_match_physical_numeric_polars_type; @@ -146,7 +146,7 @@ where T: PolarsNumericType, F: Fn(T::Native) -> T::Native + Copy, { - ca.apply_generic(|v| v.map(op)) + unary_elementwise(ca, |v| v.map(op)) } fn clip_binary(ca: &ChunkedArray, bound: &ChunkedArray, op: F) -> ChunkedArray diff --git a/crates/polars-ops/src/series/ops/cum_agg.rs b/crates/polars-ops/src/series/ops/cum_agg.rs index 03aed306f118..bd498f5088e6 100644 --- a/crates/polars-ops/src/series/ops/cum_agg.rs +++ b/crates/polars-ops/src/series/ops/cum_agg.rs @@ -1,5 +1,6 @@ use std::ops::{Add, AddAssign, Mul}; +use arity::unary_elementwise_values; use num_traits::{Bounded, One, Zero}; use polars_core::prelude::*; use polars_core::series::IsSorted; @@ -216,7 +217,7 @@ pub fn cum_count(s: &Series, reverse: bool) -> PolarsResult { let out: IdxCa = if reverse { let mut count = (s.len() - s.null_count()) as IdxSize; let mut prev = false; - ca.apply_values_generic(|v: bool| { + unary_elementwise_values(&ca, |v: bool| { if prev { count -= 1; } @@ -225,7 +226,7 @@ pub fn cum_count(s: &Series, reverse: bool) -> PolarsResult { }) } else { let mut count = 0 as IdxSize; - ca.apply_values_generic(|v: bool| { + unary_elementwise_values(&ca, |v: bool| { if v { count += 1; } diff --git a/crates/polars-ops/src/series/ops/index.rs b/crates/polars-ops/src/series/ops/index.rs index 368780566d8f..51811cf0c319 100644 --- a/crates/polars-ops/src/series/ops/index.rs +++ b/crates/polars-ops/src/series/ops/index.rs @@ -1,5 +1,6 @@ use num_traits::{Signed, Zero}; use polars_core::error::{polars_ensure, PolarsResult}; +use polars_core::prelude::arity::unary_elementwise_values; use polars_core::prelude::{ChunkedArray, DataType, IdxCa, PolarsIntegerType, Series, IDX_DTYPE}; use polars_utils::index::ToIdx; @@ -9,7 +10,7 @@ where T::Native: ToIdx, { let target_len = target_len as u64; - Ok(ca.apply_values_generic(|v| v.to_idx(target_len))) + Ok(unary_elementwise_values(ca, |v| v.to_idx(target_len))) } pub fn convert_to_unsigned_index(s: &Series, target_len: usize) -> PolarsResult { diff --git a/crates/polars-ops/src/series/ops/is_in.rs b/crates/polars-ops/src/series/ops/is_in.rs index b4fe6d54b094..7a4bdfa3f495 100644 --- a/crates/polars-ops/src/series/ops/is_in.rs +++ b/crates/polars-ops/src/series/ops/is_in.rs @@ -1,5 +1,6 @@ use std::hash::Hash; +use polars_core::prelude::arity::unary_elementwise_values; use polars_core::prelude::*; use polars_core::utils::{try_get_supertype, CustomIterTools}; use polars_core::with_match_physical_numeric_polars_type; @@ -24,9 +25,7 @@ where } }) }); - Ok(ca - .apply_values_generic(|val| set.contains(&val.to_total_ord())) - .with_name(ca.name())) + Ok(unary_elementwise_values(ca, |val| set.contains(&val.to_total_ord())).with_name(ca.name())) } fn is_in_helper<'a, T>(ca: &'a ChunkedArray, other: &Series) -> PolarsResult @@ -623,10 +622,10 @@ fn is_in_cat(ca_in: &CategoricalChunked, other: &Series) -> PolarsResult fmt, None => sniff_fmt_date(string_ca)?, }; - let ca = string_ca.apply_generic(|opt_s| { + let ca = unary_elementwise(string_ca, |opt_s| { let mut s = opt_s?; let fmt_len = fmt.len(); @@ -169,38 +170,37 @@ pub trait StringMethods: AsString { TimeUnit::Milliseconds => datetime_to_timestamp_ms, }; - let ca = string_ca - .apply_generic(|opt_s| { - let mut s = opt_s?; - let fmt_len = fmt.len(); + let ca = unary_elementwise(string_ca, |opt_s| { + let mut s = opt_s?; + let fmt_len = fmt.len(); - for i in 1..(s.len().saturating_sub(fmt_len)) { - if s.is_empty() { - return None; - } - let timestamp = if tz_aware { - DateTime::parse_from_str(s, fmt).map(|dt| func(dt.naive_utc())) - } else { - NaiveDateTime::parse_from_str(s, fmt).map(func) - }; - match timestamp { - Ok(ts) => return Some(ts), - Err(e) => { - let e: ParseErrorByteCopy = e.into(); - match e.0 { - ParseErrorKind::TooLong => { - s = &s[..s.len() - 1]; - }, - _ => { - s = &s[i..]; - }, - } - }, - } + for i in 1..(s.len().saturating_sub(fmt_len)) { + if s.is_empty() { + return None; + } + let timestamp = if tz_aware { + DateTime::parse_from_str(s, fmt).map(|dt| func(dt.naive_utc())) + } else { + NaiveDateTime::parse_from_str(s, fmt).map(func) + }; + match timestamp { + Ok(ts) => return Some(ts), + Err(e) => { + let e: ParseErrorByteCopy = e.into(); + match e.0 { + ParseErrorKind::TooLong => { + s = &s[..s.len() - 1]; + }, + _ => { + s = &s[i..]; + }, + } + }, } - None - }) - .with_name(string_ca.name()); + } + None + }) + .with_name(string_ca.name()); match (tz_aware, tz) { #[cfg(feature = "timezones")] (false, Some(tz)) => polars_ops::prelude::replace_time_zone( @@ -241,7 +241,7 @@ pub trait StringMethods: AsString { }, (string_ca.len() as f64).sqrt() as usize, ); - string_ca.apply_generic(|val| convert.eval(val?, use_cache)) + unary_elementwise(string_ca, |val| convert.eval(val?, use_cache)) } else { let mut convert = FastCachedFunc::new( |s| { @@ -250,7 +250,7 @@ pub trait StringMethods: AsString { }, (string_ca.len() as f64).sqrt() as usize, ); - string_ca.apply_generic(|val| convert.eval(val?, use_cache)) + unary_elementwise(string_ca, |val| convert.eval(val?, use_cache)) }; Ok(ca.with_name(string_ca.name()).into()) @@ -291,13 +291,14 @@ pub trait StringMethods: AsString { }, (string_ca.len() as f64).sqrt() as usize, ); - Ok(string_ca - .apply_generic(|opt_s| convert.eval(opt_s?, use_cache)) - .with_name(string_ca.name()) - .into_datetime( - tu, - Some(tz.map(|x| x.to_string()).unwrap_or("UTC".to_string())), - )) + Ok( + unary_elementwise(string_ca, |opt_s| convert.eval(opt_s?, use_cache)) + .with_name(string_ca.name()) + .into_datetime( + tu, + Some(tz.map(|x| x.to_string()).unwrap_or("UTC".to_string())), + ), + ) } #[cfg(not(feature = "timezones"))] { @@ -323,13 +324,13 @@ pub trait StringMethods: AsString { }, (string_ca.len() as f64).sqrt() as usize, ); - string_ca.apply_generic(|opt_s| convert.eval(opt_s?, use_cache)) + unary_elementwise(string_ca, |opt_s| convert.eval(opt_s?, use_cache)) } else { let mut convert = FastCachedFunc::new( |s| transform(s, &fmt), (string_ca.len() as f64).sqrt() as usize, ); - string_ca.apply_generic(|opt_s| convert.eval(opt_s?, use_cache)) + unary_elementwise(string_ca, |opt_s| convert.eval(opt_s?, use_cache)) }; let dt = ca.with_name(string_ca.name()).into_datetime(tu, None); match tz { diff --git a/crates/polars/src/docs/eager.rs b/crates/polars/src/docs/eager.rs index 28e12057bde8..8faf05c0a96c 100644 --- a/crates/polars/src/docs/eager.rs +++ b/crates/polars/src/docs/eager.rs @@ -247,6 +247,7 @@ //! //! ``` //! use polars::prelude::*; +//! use polars::prelude::arity::unary_elementwise_values; //! # fn example() -> PolarsResult<()> { //! //! // apply a closure over all values @@ -255,7 +256,7 @@ //! //! // count string lengths //! let s = Series::new("foo", &["foo", "bar", "foobar"]); -//! s.str()?.apply_values_generic(|str_val| str_val.len() as u64); +//! unary_elementwise_values(s.str()?, |str_val| str_val.len() as u64); //! //! # Ok(()) //! # }