diff --git a/crates/polars-ops/src/chunked_array/strings/namespace.rs b/crates/polars-ops/src/chunked_array/strings/namespace.rs index be8b2d71a6bd..e8e08906b84d 100644 --- a/crates/polars-ops/src/chunked_array/strings/namespace.rs +++ b/crates/polars-ops/src/chunked_array/strings/namespace.rs @@ -416,19 +416,29 @@ pub trait Utf8NameSpaceImpl: AsUtf8 { } /// Return the first n characters - fn str_head(&self, n: u64) -> PolarsResult { + fn str_head(&self, n: i64) -> PolarsResult { let ca = self.as_utf8(); - let chunks = ca + + // if n is negative, we return all but the last abs(n) characters + let chunks = if n < 0 { + let n = n.abs() as u64; + ca .downcast_iter() - .map(|c| substring(c, 0, &Some(n))) - .collect::>()?; + .map(|c| substring(c, 0, &Some(c.len() as u64 - n))) + .collect::>()? + } else { + let n = n as u64; + ca + .downcast_iter() + .map(|c| substring(c, 0, &Some(n as u64))) + .collect::>()? + }; unsafe { Ok(Utf8Chunked::from_chunks(ca.name(), chunks)) } } /// Return the last n characters - fn str_tail(&self, n: u64) -> PolarsResult { - let n = i64::try_from(n).unwrap(); + fn str_tail(&self, n: i64) -> PolarsResult { let ca = self.as_utf8(); let chunks = ca .downcast_iter() diff --git a/crates/polars-plan/src/dsl/function_expr/strings.rs b/crates/polars-plan/src/dsl/function_expr/strings.rs index c1a105bac59b..4e910bd7d2d4 100644 --- a/crates/polars-plan/src/dsl/function_expr/strings.rs +++ b/crates/polars-plan/src/dsl/function_expr/strings.rs @@ -69,8 +69,8 @@ pub enum StringFunction { }, RStrip(Option), Slice(i64, Option), - Head(u64), - Tail(u64), + Head(i64), + Tail(i64), StartsWith, Strip(Option), #[cfg(feature = "temporal")] @@ -733,12 +733,12 @@ pub(super) fn str_slice(s: &Series, start: i64, length: Option) -> PolarsRe ca.str_slice(start, length).map(|ca| ca.into_series()) } -pub(super) fn str_head(s: &Series, n: u64) -> PolarsResult { +pub(super) fn str_head(s: &Series, n: i64) -> PolarsResult { let ca = s.utf8()?; ca.str_head(n).map(|ca| ca.into_series()) } -pub(super) fn str_tail(s: &Series, n: u64) -> PolarsResult { +pub(super) fn str_tail(s: &Series, n: i64) -> PolarsResult { let ca = s.utf8()?; ca.str_tail(n).map(|ca| ca.into_series()) } diff --git a/crates/polars-plan/src/dsl/string.rs b/crates/polars-plan/src/dsl/string.rs index e47fcc1aa710..6b925f074d6b 100644 --- a/crates/polars-plan/src/dsl/string.rs +++ b/crates/polars-plan/src/dsl/string.rs @@ -517,13 +517,13 @@ impl StringNameSpace { } /// Return the first n characters in the string - pub fn str_head(self, n: u64) -> Expr { + pub fn str_head(self, n: i64) -> Expr { self.0 .map_private(FunctionExpr::StringExpr(StringFunction::Head(n))) } /// Return the last n characters in the string - pub fn str_tail(self, n: u64) -> Expr { + pub fn str_tail(self, n: i64) -> Expr { self.0 .map_private(FunctionExpr::StringExpr(StringFunction::Tail(n))) } diff --git a/py-polars/src/expr/string.rs b/py-polars/src/expr/string.rs index d8ac30b1123d..f06b2b80a619 100644 --- a/py-polars/src/expr/string.rs +++ b/py-polars/src/expr/string.rs @@ -73,11 +73,11 @@ impl PyExpr { self.inner.clone().str().str_slice(start, length).into() } - fn str_head(&self, n: u64) -> Self { + fn str_head(&self, n: i64) -> Self { self.inner.clone().str().str_head(n).into() } - fn str_tail(&self, n: u64) -> Self { + fn str_tail(&self, n: i64) -> Self { self.inner.clone().str().str_tail(n).into() }