Skip to main content

ferray_strings/
str_ops.rs

1// ferray-strings: Miscellaneous string operations
2//
3// str_len, swapcase, and elementwise comparison functions.
4
5use ferray_core::Array;
6use ferray_core::dimension::Dimension;
7use ferray_core::error::FerrayResult;
8
9use crate::string_array::StringArray;
10
11/// Return the length of each string element. Matches
12/// `numpy.strings.str_len` (#518).
13pub fn str_len<D: Dimension>(a: &StringArray<D>) -> FerrayResult<Array<u64, D>> {
14    let data: Vec<u64> = a.iter().map(|s| s.len() as u64).collect();
15    Array::from_vec(a.dim().clone(), data)
16}
17
18/// Swap the case of each character in every element. Matches
19/// `numpy.strings.swapcase` (#515).
20pub fn swapcase<D: Dimension>(a: &StringArray<D>) -> FerrayResult<StringArray<D>> {
21    a.map(|s| {
22        s.chars()
23            .map(|c| {
24                if c.is_uppercase() {
25                    c.to_lowercase().collect::<String>()
26                } else if c.is_lowercase() {
27                    c.to_uppercase().collect::<String>()
28                } else {
29                    c.to_string()
30                }
31            })
32            .collect()
33    })
34}
35
36// ---------------------------------------------------------------------------
37// Elementwise string comparison (#516)
38// ---------------------------------------------------------------------------
39
40/// Elementwise string equality. Both arrays must have the same shape.
41pub fn equal<D: Dimension>(a: &StringArray<D>, b: &StringArray<D>) -> FerrayResult<Array<bool, D>> {
42    let data: Vec<bool> = a.iter().zip(b.iter()).map(|(x, y)| x == y).collect();
43    Array::from_vec(a.dim().clone(), data)
44}
45
46/// Elementwise string inequality.
47pub fn not_equal<D: Dimension>(
48    a: &StringArray<D>,
49    b: &StringArray<D>,
50) -> FerrayResult<Array<bool, D>> {
51    let data: Vec<bool> = a.iter().zip(b.iter()).map(|(x, y)| x != y).collect();
52    Array::from_vec(a.dim().clone(), data)
53}
54
55/// Elementwise lexicographic less-than.
56pub fn less<D: Dimension>(a: &StringArray<D>, b: &StringArray<D>) -> FerrayResult<Array<bool, D>> {
57    let data: Vec<bool> = a.iter().zip(b.iter()).map(|(x, y)| x < y).collect();
58    Array::from_vec(a.dim().clone(), data)
59}
60
61/// Elementwise lexicographic greater-than.
62pub fn greater<D: Dimension>(
63    a: &StringArray<D>,
64    b: &StringArray<D>,
65) -> FerrayResult<Array<bool, D>> {
66    let data: Vec<bool> = a.iter().zip(b.iter()).map(|(x, y)| x > y).collect();
67    Array::from_vec(a.dim().clone(), data)
68}
69
70/// Elementwise lexicographic less-or-equal.
71pub fn less_equal<D: Dimension>(
72    a: &StringArray<D>,
73    b: &StringArray<D>,
74) -> FerrayResult<Array<bool, D>> {
75    let data: Vec<bool> = a.iter().zip(b.iter()).map(|(x, y)| x <= y).collect();
76    Array::from_vec(a.dim().clone(), data)
77}
78
79/// Elementwise lexicographic greater-or-equal.
80pub fn greater_equal<D: Dimension>(
81    a: &StringArray<D>,
82    b: &StringArray<D>,
83) -> FerrayResult<Array<bool, D>> {
84    let data: Vec<bool> = a.iter().zip(b.iter()).map(|(x, y)| x >= y).collect();
85    Array::from_vec(a.dim().clone(), data)
86}
87
88#[cfg(test)]
89mod tests {
90    use super::*;
91    use crate::string_array::array;
92
93    #[test]
94    fn test_str_len() {
95        let a = array(&["hello", "", "abc", "hi"]).unwrap();
96        let r = str_len(&a).unwrap();
97        assert_eq!(r.as_slice().unwrap(), &[5, 0, 3, 2]);
98    }
99
100    #[test]
101    fn test_swapcase() {
102        let a = array(&["Hello World", "ABC", "abc", "123"]).unwrap();
103        let r = swapcase(&a).unwrap();
104        assert_eq!(r.as_slice(), &["hELLO wORLD", "abc", "ABC", "123"]);
105    }
106
107    #[test]
108    fn test_equal() {
109        let a = array(&["abc", "def", "ghi"]).unwrap();
110        let b = array(&["abc", "xyz", "ghi"]).unwrap();
111        let r = equal(&a, &b).unwrap();
112        assert_eq!(r.as_slice().unwrap(), &[true, false, true]);
113    }
114
115    #[test]
116    fn test_less() {
117        let a = array(&["abc", "xyz"]).unwrap();
118        let b = array(&["abd", "abc"]).unwrap();
119        let r = less(&a, &b).unwrap();
120        assert_eq!(r.as_slice().unwrap(), &[true, false]);
121    }
122
123    #[test]
124    fn test_greater() {
125        let a = array(&["xyz", "abc"]).unwrap();
126        let b = array(&["abc", "xyz"]).unwrap();
127        let r = greater(&a, &b).unwrap();
128        assert_eq!(r.as_slice().unwrap(), &[true, false]);
129    }
130}