polars_ops/chunked_array/binary/
namespace.rs

1#[cfg(any(
2    feature = "binary_encoding",
3    all(feature = "strings", feature = "string_encoding")
4))]
5use base64::Engine as _;
6#[cfg(any(
7    feature = "binary_encoding",
8    all(feature = "strings", feature = "string_encoding")
9))]
10use base64::engine::general_purpose;
11use memchr::memmem::find;
12#[cfg(feature = "binary_encoding")]
13use polars_arrow::array::{Array, MutablePlBinary};
14use polars_compute::cast::{binview_to_fixed_size_list_dyn, binview_to_primitive_dyn};
15use polars_compute::size::binary_size_bytes;
16use polars_core::prelude::arity::{broadcast_binary_elementwise_values, unary_elementwise_values};
17
18use super::*;
19
20#[cfg(any(
21    feature = "binary_encoding",
22    all(feature = "strings", feature = "string_encoding")
23))]
24pub(crate) fn binary_to_hex(ca: &BinaryChunked) -> StringChunked {
25    ca.apply_into_string_amortized(|s, buf| {
26        // SAFETY: hex output is ASCII.
27        let out = unsafe { buf.as_mut_vec() };
28        out.resize(s.len() * 2, 0);
29        hex::encode_to_slice(s, out).unwrap();
30    })
31}
32
33#[cfg(any(
34    feature = "binary_encoding",
35    all(feature = "strings", feature = "string_encoding")
36))]
37pub(crate) fn binary_to_base64(ca: &BinaryChunked) -> StringChunked {
38    ca.apply_into_string_amortized(|s, buf| general_purpose::STANDARD.encode_string(s, buf))
39}
40
41/// Decode every value into a reused buffer. `decode` returns `false` on invalid input.
42#[cfg(feature = "binary_encoding")]
43fn decode_amortized(
44    ca: &BinaryChunked,
45    strict: bool,
46    encoding: &str,
47    mut decode: impl FnMut(&[u8], &mut Vec<u8>) -> bool,
48) -> PolarsResult<BinaryChunked> {
49    let mut buf = Vec::new();
50    let chunks = ca.downcast_iter().map(|arr| {
51        let mut out = MutablePlBinary::with_capacity(arr.len());
52        for opt_s in arr.iter() {
53            match opt_s {
54                None => out.push_null(),
55                Some(s) => {
56                    buf.clear();
57                    if decode(s, &mut buf) {
58                        out.push_value(&buf);
59                    } else if strict {
60                        polars_bail!(
61                            ComputeError:
62                            "invalid `{encoding}` encoding found; try setting `strict=false` to ignore"
63                        );
64                    } else {
65                        out.push_null();
66                    }
67                },
68            }
69        }
70        Ok(out.freeze())
71    });
72    ChunkedArray::try_from_chunk_iter(ca.name().clone(), chunks)
73}
74
75pub trait BinaryNameSpaceImpl: AsBinary {
76    /// Slice the binary values.
77    ///
78    /// Determines a slice starting from `offset` and with length `length` of each of the elements.
79    /// `offset` can be negative, in which case the start counts from the end of the bytes.
80    fn bin_slice(&self, offset: &Column, length: &Column) -> PolarsResult<BinaryChunked> {
81        let ca = self.as_binary();
82        let offset = offset.cast(&DataType::Int64)?;
83        let length = length.strict_cast(&DataType::UInt64)?;
84
85        Ok(super::slice::slice(ca, offset.i64()?, length.u64()?))
86    }
87    /// Slice the first `n` bytes of the binary value.
88    ///
89    /// Determines a slice starting at the beginning of the binary data up to offset `n` of each
90    /// element. `n` can be negative, in which case the slice ends `n` bytes from the end.
91    fn bin_head(&self, n: &Column) -> PolarsResult<BinaryChunked> {
92        let ca = self.as_binary();
93        let n = n.strict_cast(&DataType::Int64)?;
94
95        super::slice::head(ca, n.i64()?)
96    }
97
98    /// Slice the last `n` bytes of the binary value.
99    ///
100    /// Determines a slice starting at offset `n` of each element. `n` can be
101    /// negative, in which case the slice begins `n` bytes from the start.
102    fn bin_tail(&self, n: &Column) -> PolarsResult<BinaryChunked> {
103        let ca = self.as_binary();
104        let n = n.strict_cast(&DataType::Int64)?;
105
106        super::slice::tail(ca, n.i64()?)
107    }
108
109    /// Check if binary contains given literal
110    fn contains(&self, lit: &[u8]) -> BooleanChunked {
111        let ca = self.as_binary();
112        let f = |s: &[u8]| find(s, lit).is_some();
113        unary_elementwise_values(ca, f)
114    }
115
116    fn contains_chunked(&self, lit: &BinaryChunked) -> PolarsResult<BooleanChunked> {
117        let ca = self.as_binary();
118        Ok(match lit.len() {
119            1 => match lit.get(0) {
120                Some(lit) => ca.contains(lit),
121                None => BooleanChunked::full_null(ca.name().clone(), ca.len()),
122            },
123            _ => {
124                polars_ensure!(
125                    ca.len() == lit.len() || ca.len() == 1,
126                    length_mismatch = "bin.contains",
127                    ca.len(),
128                    lit.len()
129                );
130                broadcast_binary_elementwise_values(ca, lit, |src, lit| find(src, lit).is_some())
131            },
132        })
133    }
134
135    /// Check if strings ends with a substring
136    fn ends_with(&self, sub: &[u8]) -> BooleanChunked {
137        let ca = self.as_binary();
138        let f = |s: &[u8]| s.ends_with(sub);
139        ca.apply_nonnull_values_generic(DataType::Boolean, f)
140    }
141
142    /// Check if strings starts with a substring
143    fn starts_with(&self, sub: &[u8]) -> BooleanChunked {
144        let ca = self.as_binary();
145        let f = |s: &[u8]| s.starts_with(sub);
146        ca.apply_nonnull_values_generic(DataType::Boolean, f)
147    }
148
149    fn starts_with_chunked(&self, prefix: &BinaryChunked) -> PolarsResult<BooleanChunked> {
150        let ca = self.as_binary();
151        Ok(match prefix.len() {
152            1 => match prefix.get(0) {
153                Some(s) => self.starts_with(s),
154                None => BooleanChunked::full_null(ca.name().clone(), ca.len()),
155            },
156            _ => {
157                polars_ensure!(
158                    ca.len() == prefix.len() || ca.len() == 1,
159                    length_mismatch = "bin.starts_with",
160                    ca.len(),
161                    prefix.len()
162                );
163                broadcast_binary_elementwise_values(ca, prefix, |s, sub| s.starts_with(sub))
164            },
165        })
166    }
167
168    fn ends_with_chunked(&self, suffix: &BinaryChunked) -> PolarsResult<BooleanChunked> {
169        let ca = self.as_binary();
170        Ok(match suffix.len() {
171            1 => match suffix.get(0) {
172                Some(s) => self.ends_with(s),
173                None => BooleanChunked::full_null(ca.name().clone(), ca.len()),
174            },
175            _ => {
176                polars_ensure!(
177                    ca.len() == suffix.len() || ca.len() == 1,
178                    length_mismatch = "bin.ends_with",
179                    ca.len(),
180                    suffix.len()
181                );
182                broadcast_binary_elementwise_values(ca, suffix, |s, sub| s.ends_with(sub))
183            },
184        })
185    }
186
187    /// Get the size of the binary values in bytes.
188    fn size_bytes(&self) -> UInt32Chunked {
189        let ca = self.as_binary();
190        ca.apply_kernel_cast(&binary_size_bytes)
191    }
192
193    #[cfg(feature = "binary_encoding")]
194    fn hex_decode(&self, strict: bool) -> PolarsResult<BinaryChunked> {
195        decode_amortized(self.as_binary(), strict, "hex", |s, buf| {
196            buf.resize(s.len() / 2, 0);
197            hex::decode_to_slice(s, buf).is_ok()
198        })
199    }
200
201    #[cfg(feature = "binary_encoding")]
202    fn hex_encode(&self) -> Series {
203        binary_to_hex(self.as_binary()).into_series()
204    }
205
206    #[cfg(feature = "binary_encoding")]
207    fn base64_decode(&self, strict: bool) -> PolarsResult<BinaryChunked> {
208        decode_amortized(self.as_binary(), strict, "base64", |s, buf| {
209            general_purpose::STANDARD.decode_vec(s, buf).is_ok()
210        })
211    }
212
213    #[cfg(feature = "binary_encoding")]
214    fn base64_encode(&self) -> Series {
215        binary_to_base64(self.as_binary()).into_series()
216    }
217
218    #[cfg(feature = "binary_encoding")]
219    fn reinterpret(&self, dtype: &DataType, is_little_endian: bool) -> PolarsResult<Series> {
220        unsafe {
221            Ok(Series::from_chunks_and_dtype_unchecked(
222                self.as_binary().name().clone(),
223                self._reinterpret_inner(dtype, is_little_endian)?,
224                dtype,
225            ))
226        }
227    }
228
229    #[cfg(feature = "binary_encoding")]
230    fn _reinterpret_inner(
231        &self,
232        dtype: &DataType,
233        is_little_endian: bool,
234    ) -> PolarsResult<Vec<Box<dyn Array>>> {
235        use polars_core::with_match_physical_numeric_polars_type;
236
237        let ca = self.as_binary();
238
239        match dtype {
240            dtype if dtype.is_primitive_numeric() || dtype.is_temporal() => {
241                let dtype = dtype.to_physical();
242                let arrow_data_type = dtype
243                    .to_arrow(CompatLevel::newest())
244                    .underlying_physical_type();
245                with_match_physical_numeric_polars_type!(dtype, |$T| {
246                    unsafe {
247                        ca.chunks().iter().map(|chunk| {
248                            binview_to_primitive_dyn::<<$T as PolarsNumericType>::Native>(
249                                &**chunk,
250                                &arrow_data_type,
251                                is_little_endian,
252                            )
253                        }).collect()
254                    }
255                })
256            },
257            #[cfg(feature = "dtype-array")]
258            DataType::Array(inner_dtype, array_width)
259                if inner_dtype.is_primitive_numeric() || inner_dtype.is_temporal() =>
260            {
261                let inner_dtype = inner_dtype.to_physical();
262                let result: Vec<ArrayRef> = with_match_physical_numeric_polars_type!(inner_dtype, |$T| {
263                    unsafe {
264                        ca.chunks().iter().map(|chunk| {
265                            binview_to_fixed_size_list_dyn::<<$T as PolarsNumericType>::Native>(
266                                &**chunk,
267                                *array_width,
268                                is_little_endian
269                            )
270                        }).collect::<Result<Vec<ArrayRef>, _>>()
271                    }
272                })?;
273                Ok(result)
274            },
275            _ => Err(
276                polars_err!(InvalidOperation: "unsupported data type {:?} in reinterpret. Only numeric or temporal types, or Arrays of those, are allowed.", dtype),
277            ),
278        }
279    }
280}
281
282impl BinaryNameSpaceImpl for BinaryChunked {}
polars_ops/chunked_array/binary/namespace.rs

polars_ops/chunked_array/binary/
namespace.rs
