Initial add of negative n

pola-rs · Aug 8, 2023 · 36adb83 · 36adb83
1 parent 1924999
commit 36adb83
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 14 deletions.
diff --git 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<Utf8Chunked> {
+    fn str_head(&self, n: i64) -> PolarsResult<Utf8Chunked> {
         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::<arrow::error::Result<_>>()?;
+            .map(|c| substring(c, 0, &Some(c.len() as u64 - n)))
+            .collect::<arrow::error::Result<_>>()?
+        } else {
+            let n = n as u64;
+            ca
+                .downcast_iter()
+                .map(|c| substring(c, 0, &Some(n as u64)))
+                .collect::<arrow::error::Result<_>>()?
+        };
 
         unsafe { Ok(Utf8Chunked::from_chunks(ca.name(), chunks)) }
     }
 
     /// Return the last n characters
-    fn str_tail(&self, n: u64) -> PolarsResult<Utf8Chunked> {
-        let n = i64::try_from(n).unwrap();
+    fn str_tail(&self, n: i64) -> PolarsResult<Utf8Chunked> {
         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
@@ -69,8 +69,8 @@ pub enum StringFunction {
     },
     RStrip(Option<String>),
     Slice(i64, Option<u64>),
-    Head(u64),
-    Tail(u64),
+    Head(i64),
+    Tail(i64),
     StartsWith,
     Strip(Option<String>),
     #[cfg(feature = "temporal")]
@@ -733,12 +733,12 @@ pub(super) fn str_slice(s: &Series, start: i64, length: Option<u64>) -> PolarsRe
     ca.str_slice(start, length).map(|ca| ca.into_series())
 }
 
-pub(super) fn str_head(s: &Series, n: u64) -> PolarsResult<Series> {
+pub(super) fn str_head(s: &Series, n: i64) -> PolarsResult<Series> {
     let ca = s.utf8()?;
     ca.str_head(n).map(|ca| ca.into_series())
 }
 
-pub(super) fn str_tail(s: &Series, n: u64) -> PolarsResult<Series> {
+pub(super) fn str_tail(s: &Series, n: i64) -> PolarsResult<Series> {
     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
@@ -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
@@ -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()
     }