datafusion_functions/unicode/
common.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
18//! Common utilities for implementing unicode functions
19
20use arrow::array::{
21    Array, ArrayRef, ByteView, GenericStringArray, Int64Array, OffsetSizeTrait,
22    StringViewArray, make_view,
23};
24use arrow::datatypes::DataType;
25use arrow_buffer::{NullBuffer, ScalarBuffer};
26use datafusion_common::Result;
27use datafusion_common::ScalarValue;
28use datafusion_common::cast::{
29    as_generic_string_array, as_int64_array, as_string_view_array,
30};
31use datafusion_common::exec_err;
32use datafusion_expr::ColumnarValue;
33use std::cmp::Ordering;
34use std::ops::Range;
35use std::sync::Arc;
36
37/// If `cv` is a non-null scalar string, return its value.
38pub(crate) fn try_as_scalar_str(cv: &ColumnarValue) -> Option<&str> {
39    match cv {
40        ColumnarValue::Scalar(s) => s.try_as_str().flatten(),
41        _ => None,
42    }
43}
44
45/// If `cv` is a non-null scalar Int64, return its value.
46pub(crate) fn try_as_scalar_i64(cv: &ColumnarValue) -> Option<i64> {
47    match cv {
48        ColumnarValue::Scalar(ScalarValue::Int64(v)) => *v,
49        _ => None,
50    }
51}
52
53/// Estimates data capacity for `pad` based on `length_array` with row length.
54/// For ASCII, one row is at most `target_len` bytes.
55/// For UTF8, it could be larger
56pub(crate) fn pad_data_capacity(length_array: &Int64Array) -> usize {
57    length_array
58        .iter()
59        .flatten()
60        .fold(0, |acc, len| acc.saturating_add(len as usize))
61}
62
63/// A trait for `left` and `right` byte slicing operations
64pub(crate) trait LeftRightSlicer {
65    fn slice(string: &str, n: i64) -> Range<usize>;
66}
67
68pub(crate) struct LeftSlicer {}
69
70impl LeftRightSlicer for LeftSlicer {
71    fn slice(string: &str, n: i64) -> Range<usize> {
72        0..left_right_byte_length(string, n)
73    }
74}
75
76pub(crate) struct RightSlicer {}
77
78impl LeftRightSlicer for RightSlicer {
79    fn slice(string: &str, n: i64) -> Range<usize> {
80        if n == 0 {
81            // Return nothing for `n=0`
82            0..0
83        } else if n == i64::MIN {
84            // Special case for i64::MIN overflow
85            0..0
86        } else {
87            left_right_byte_length(string, -n)..string.len()
88        }
89    }
90}
91
92/// Returns the byte offset of the `n`th codepoint in `string`,
93/// or `string.len()` if the string has fewer than `n` codepoints.
94#[inline]
95pub(crate) fn byte_offset_of_char(string: &str, n: usize) -> usize {
96    string
97        .char_indices()
98        .nth(n)
99        .map_or(string.len(), |(i, _)| i)
100}
101
102/// If `string` has more than `n` codepoints, returns the byte offset of
103/// the `n`-th codepoint boundary. Otherwise returns the total codepoint count.
104#[inline]
105pub(crate) fn char_count_or_boundary(string: &str, n: usize) -> StringCharLen {
106    let mut count = 0;
107    for (byte_idx, _) in string.char_indices() {
108        if count == n {
109            return StringCharLen::ByteOffset(byte_idx);
110        }
111        count += 1;
112    }
113    StringCharLen::CharCount(count)
114}
115
116/// Result of [`char_count_or_boundary`].
117pub(crate) enum StringCharLen {
118    /// The string has more than `n` codepoints; contains the byte offset
119    /// at the `n`-th codepoint boundary.
120    ByteOffset(usize),
121    /// The string has `n` or fewer codepoints; contains the exact count.
122    CharCount(usize),
123}
124
125/// Calculate the byte length of the substring of `n` chars from string `string`
126#[inline]
127fn left_right_byte_length(string: &str, n: i64) -> usize {
128    let abs = n.unsigned_abs().min(usize::MAX as u64) as usize;
129    // For ASCII input every character is exactly one byte, so the byte offset of
130    // the n-th codepoint is just the (clamped) character count. This avoids the
131    // per-character `char_indices()` scan of the general path.
132    match n.cmp(&0) {
133        Ordering::Equal => 0,
134        // `abs` chars trimmed from the end: keep the leading `len - abs`.
135        Ordering::Less if string.is_ascii() => string.len().saturating_sub(abs),
136        Ordering::Less => string
137            .char_indices()
138            .nth_back(abs - 1)
139            .map(|(index, _)| index)
140            .unwrap_or(0),
141        // First `abs` chars, but never past the end of the string.
142        Ordering::Greater if string.is_ascii() => abs.min(string.len()),
143        Ordering::Greater => byte_offset_of_char(string, abs),
144    }
145}
146
147/// General implementation for `left` and `right` functions
148pub(crate) fn general_left_right<F: LeftRightSlicer>(
149    args: &[ArrayRef],
150) -> Result<ArrayRef> {
151    let n_array = as_int64_array(&args[1])?;
152
153    match args[0].data_type() {
154        DataType::Utf8 => {
155            let string_array = as_generic_string_array::<i32>(&args[0])?;
156            general_left_right_array::<i32, F>(string_array, n_array)
157        }
158        DataType::LargeUtf8 => {
159            let string_array = as_generic_string_array::<i64>(&args[0])?;
160            general_left_right_array::<i64, F>(string_array, n_array)
161        }
162        DataType::Utf8View => {
163            let string_view_array = as_string_view_array(&args[0])?;
164            general_left_right_view::<F>(string_view_array, n_array)
165        }
166        _ => exec_err!("Not supported"),
167    }
168}
169
170/// `left`/`right` for Utf8/LargeUtf8 input.
171fn general_left_right_array<T: OffsetSizeTrait, F: LeftRightSlicer>(
172    string_array: &GenericStringArray<T>,
173    n_array: &Int64Array,
174) -> Result<ArrayRef> {
175    let result = string_array
176        .iter()
177        .zip(n_array.iter())
178        .map(|(string, n)| match (string, n) {
179            (Some(string), Some(n)) => Some(&string[F::slice(string, n)]),
180            _ => None,
181        })
182        .collect::<GenericStringArray<T>>();
183    Ok(Arc::new(result) as ArrayRef)
184}
185
186/// `general_left_right` for StringViewArray input.
187fn general_left_right_view<F: LeftRightSlicer>(
188    string_view_array: &StringViewArray,
189    n_array: &Int64Array,
190) -> Result<ArrayRef> {
191    let views = string_view_array.views();
192    let new_nulls = NullBuffer::union(string_view_array.nulls(), n_array.nulls());
193    let len = n_array.len();
194    let mut has_out_of_line = false;
195
196    let new_views = (0..len)
197        .map(|idx| {
198            if new_nulls.as_ref().is_some_and(|n| n.is_null(idx)) {
199                return 0;
200            }
201
202            // SAFETY: we just checked validity above
203            let string: &str = unsafe { string_view_array.value_unchecked(idx) };
204            let n = n_array.value(idx);
205
206            let range = F::slice(string, n);
207            let result_bytes = &string.as_bytes()[range.clone()];
208            if result_bytes.len() > 12 {
209                has_out_of_line = true;
210            }
211
212            let byte_view = ByteView::from(views[idx]);
213            let new_offset = byte_view.offset + (range.start as u32);
214            make_view(result_bytes, byte_view.buffer_index, new_offset)
215        })
216        .collect::<Vec<u128>>();
217
218    let views = ScalarBuffer::from(new_views);
219    let data_buffers = if has_out_of_line {
220        string_view_array.data_buffers().to_vec()
221    } else {
222        vec![]
223    };
224
225    // SAFETY:
226    // - Each view is produced by `make_view` with correct bytes and offset
227    // - Out-of-line views reuse the original buffer index and adjusted offset
228    unsafe {
229        let array = StringViewArray::new_unchecked(views, data_buffers, new_nulls);
230        Ok(Arc::new(array) as ArrayRef)
231    }
232}
datafusion_functions/unicode/common.rs

datafusion_functions/unicode/
common.rs
