arrow_string/
binary_predicate.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use arrow_array::{Array, ArrayAccessor, BinaryViewArray, BooleanArray};
19use arrow_buffer::BooleanBuffer;
20use memchr::memmem::Finder;
21use std::iter::zip;
22
23/// A binary based predicate
24pub enum BinaryPredicate<'a> {
25    Contains(Finder<'a>),
26    StartsWith(&'a [u8]),
27    EndsWith(&'a [u8]),
28}
29
30impl<'a> BinaryPredicate<'a> {
31    pub fn contains(needle: &'a [u8]) -> Self {
32        Self::Contains(Finder::new(needle))
33    }
34
35    /// Evaluate this predicate against the given haystack
36    pub fn evaluate(&self, haystack: &[u8]) -> bool {
37        match self {
38            Self::Contains(finder) => finder.find(haystack).is_some(),
39            Self::StartsWith(v) => starts_with(haystack, v, equals_kernel),
40            Self::EndsWith(v) => ends_with(haystack, v, equals_kernel),
41        }
42    }
43
44    /// Evaluate this predicate against the elements of `array`
45    ///
46    /// If `negate` is true the result of the predicate will be negated
47    #[inline(never)]
48    pub fn evaluate_array<'i, T>(&self, array: T, negate: bool) -> BooleanArray
49    where
50        T: ArrayAccessor<Item = &'i [u8]>,
51    {
52        match self {
53            Self::Contains(finder) => BooleanArray::from_unary(array, |haystack| {
54                finder.find(haystack).is_some() != negate
55            }),
56            Self::StartsWith(v) => {
57                if let Some(view_array) = array.as_any().downcast_ref::<BinaryViewArray>() {
58                    let nulls = view_array.logical_nulls();
59                    let values = BooleanBuffer::from(
60                        view_array
61                            .prefix_bytes_iter(v.len())
62                            .map(|haystack| equals_bytes(haystack, v, equals_kernel) != negate)
63                            .collect::<Vec<_>>(),
64                    );
65                    BooleanArray::new(values, nulls)
66                } else {
67                    BooleanArray::from_unary(array, |haystack| {
68                        starts_with(haystack, v, equals_kernel) != negate
69                    })
70                }
71            }
72            Self::EndsWith(v) => {
73                if let Some(view_array) = array.as_any().downcast_ref::<BinaryViewArray>() {
74                    let nulls = view_array.logical_nulls();
75                    let values = BooleanBuffer::from(
76                        view_array
77                            .suffix_bytes_iter(v.len())
78                            .map(|haystack| equals_bytes(haystack, v, equals_kernel) != negate)
79                            .collect::<Vec<_>>(),
80                    );
81                    BooleanArray::new(values, nulls)
82                } else {
83                    BooleanArray::from_unary(array, |haystack| {
84                        ends_with(haystack, v, equals_kernel) != negate
85                    })
86                }
87            }
88        }
89    }
90}
91
92fn equals_bytes(lhs: &[u8], rhs: &[u8], byte_eq_kernel: impl Fn((&u8, &u8)) -> bool) -> bool {
93    lhs.len() == rhs.len() && zip(lhs, rhs).all(byte_eq_kernel)
94}
95
96/// This is faster than `[u8]::starts_with` for small slices.
97/// See <https://github.com/apache/arrow-rs/issues/6107> for more details.
98fn starts_with(
99    haystack: &[u8],
100    needle: &[u8],
101    byte_eq_kernel: impl Fn((&u8, &u8)) -> bool,
102) -> bool {
103    if needle.len() > haystack.len() {
104        false
105    } else {
106        zip(haystack, needle).all(byte_eq_kernel)
107    }
108}
109/// This is faster than `[u8]::ends_with` for small slices.
110/// See <https://github.com/apache/arrow-rs/issues/6107> for more details.
111fn ends_with(haystack: &[u8], needle: &[u8], byte_eq_kernel: impl Fn((&u8, &u8)) -> bool) -> bool {
112    if needle.len() > haystack.len() {
113        false
114    } else {
115        zip(haystack.iter().rev(), needle.iter().rev()).all(byte_eq_kernel)
116    }
117}
118
119fn equals_kernel((n, h): (&u8, &u8)) -> bool {
120    n == h
121}
122
123#[cfg(test)]
124mod tests {
125    use super::BinaryPredicate;
126
127    #[test]
128    fn test_contains() {
129        assert!(BinaryPredicate::contains(b"hay").evaluate(b"haystack"));
130        assert!(BinaryPredicate::contains(b"haystack").evaluate(b"haystack"));
131        assert!(BinaryPredicate::contains(b"h").evaluate(b"haystack"));
132        assert!(BinaryPredicate::contains(b"k").evaluate(b"haystack"));
133        assert!(BinaryPredicate::contains(b"stack").evaluate(b"haystack"));
134        assert!(BinaryPredicate::contains(b"sta").evaluate(b"haystack"));
135        assert!(BinaryPredicate::contains(b"stack").evaluate(b"hay\0stack"));
136        assert!(BinaryPredicate::contains(b"\0s").evaluate(b"hay\0stack"));
137        assert!(BinaryPredicate::contains(b"\0").evaluate(b"hay\0stack"));
138        assert!(BinaryPredicate::contains(b"a").evaluate(b"a"));
139        // not matching
140        assert!(!BinaryPredicate::contains(b"hy").evaluate(b"haystack"));
141        assert!(!BinaryPredicate::contains(b"stackx").evaluate(b"haystack"));
142        assert!(!BinaryPredicate::contains(b"x").evaluate(b"haystack"));
143        assert!(!BinaryPredicate::contains(b"haystack haystack").evaluate(b"haystack"));
144    }
145
146    #[test]
147    fn test_starts_with() {
148        assert!(BinaryPredicate::StartsWith(b"hay").evaluate(b"haystack"));
149        assert!(BinaryPredicate::StartsWith(b"h\0ay").evaluate(b"h\0aystack"));
150        assert!(BinaryPredicate::StartsWith(b"haystack").evaluate(b"haystack"));
151        assert!(BinaryPredicate::StartsWith(b"ha").evaluate(b"haystack"));
152        assert!(BinaryPredicate::StartsWith(b"h").evaluate(b"haystack"));
153        assert!(BinaryPredicate::StartsWith(b"").evaluate(b"haystack"));
154
155        assert!(!BinaryPredicate::StartsWith(b"stack").evaluate(b"haystack"));
156        assert!(!BinaryPredicate::StartsWith(b"haystacks").evaluate(b"haystack"));
157        assert!(!BinaryPredicate::StartsWith(b"HAY").evaluate(b"haystack"));
158        assert!(!BinaryPredicate::StartsWith(b"h\0ay").evaluate(b"haystack"));
159        assert!(!BinaryPredicate::StartsWith(b"hay").evaluate(b"h\0aystack"));
160    }
161
162    #[test]
163    fn test_ends_with() {
164        assert!(BinaryPredicate::EndsWith(b"stack").evaluate(b"haystack"));
165        assert!(BinaryPredicate::EndsWith(b"st\0ack").evaluate(b"hayst\0ack"));
166        assert!(BinaryPredicate::EndsWith(b"haystack").evaluate(b"haystack"));
167        assert!(BinaryPredicate::EndsWith(b"ck").evaluate(b"haystack"));
168        assert!(BinaryPredicate::EndsWith(b"k").evaluate(b"haystack"));
169        assert!(BinaryPredicate::EndsWith(b"").evaluate(b"haystack"));
170
171        assert!(!BinaryPredicate::EndsWith(b"hay").evaluate(b"haystack"));
172        assert!(!BinaryPredicate::EndsWith(b"STACK").evaluate(b"haystack"));
173        assert!(!BinaryPredicate::EndsWith(b"haystacks").evaluate(b"haystack"));
174        assert!(!BinaryPredicate::EndsWith(b"xhaystack").evaluate(b"haystack"));
175        assert!(!BinaryPredicate::EndsWith(b"st\0ack").evaluate(b"haystack"));
176        assert!(!BinaryPredicate::EndsWith(b"stack").evaluate(b"hayst\0ack"));
177    }
178}
arrow_string/binary_predicate.rs

arrow_string/
binary_predicate.rs
