arrow_cast/cast/
string.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 crate::cast::*;
19use arrow_buffer::NullBuffer;
20
21pub(crate) fn cast_bool_to_string<A>(array: &dyn Array) -> Result<ArrayRef, ArrowError>
22where
23    A: Array + FromIterator<Option<&'static str>> + 'static,
24{
25    let array: A = array
26        .as_boolean()
27        .iter()
28        .map(|v| v.map(|b| if b { "true" } else { "false" }))
29        .collect();
30    Ok(Arc::new(array))
31}
32
33pub(crate) fn value_to_string<O: OffsetSizeTrait>(
34    array: &dyn Array,
35    options: &CastOptions,
36) -> Result<ArrayRef, ArrowError> {
37    let mut builder = GenericStringBuilder::<O>::new();
38    let formatter = ArrayFormatter::try_new(array, &options.format_options)?;
39    let nulls = array.nulls();
40    for i in 0..array.len() {
41        match nulls.is_some_and(|x| x.is_null(i)) {
42            true => builder.append_null(),
43            false => {
44                formatter.value(i).write(&mut builder)?;
45                // tell the builder the row is finished
46                builder.append_value("");
47            }
48        }
49    }
50    Ok(Arc::new(builder.finish()))
51}
52
53pub(crate) fn value_to_string_view(
54    array: &dyn Array,
55    options: &CastOptions,
56) -> Result<ArrayRef, ArrowError> {
57    let mut builder = StringViewBuilder::with_capacity(array.len());
58    let formatter = ArrayFormatter::try_new(array, &options.format_options)?;
59    let nulls = array.nulls();
60    // buffer to avoid reallocating on each value
61    // TODO: replace with write to builder after https://github.com/apache/arrow-rs/issues/6373
62    let mut buffer = String::new();
63    for i in 0..array.len() {
64        match nulls.is_some_and(|x| x.is_null(i)) {
65            true => builder.append_null(),
66            false => {
67                // write to buffer first and then copy into target array
68                buffer.clear();
69                formatter.value(i).write(&mut buffer)?;
70                builder.append_value(&buffer)
71            }
72        }
73    }
74    Ok(Arc::new(builder.finish()))
75}
76
77/// Parse UTF-8
78pub(crate) fn parse_string<P: Parser, O: OffsetSizeTrait>(
79    array: &dyn Array,
80    cast_options: &CastOptions,
81) -> Result<ArrayRef, ArrowError> {
82    let string_array = array.as_string::<O>();
83    parse_string_iter::<P, _, _>(string_array.iter(), cast_options, || {
84        string_array.nulls().cloned()
85    })
86}
87
88/// Parse UTF-8 View
89pub(crate) fn parse_string_view<P: Parser>(
90    array: &dyn Array,
91    cast_options: &CastOptions,
92) -> Result<ArrayRef, ArrowError> {
93    let string_view_array = array.as_string_view();
94    parse_string_iter::<P, _, _>(string_view_array.iter(), cast_options, || {
95        string_view_array.nulls().cloned()
96    })
97}
98
99fn parse_string_iter<
100    'a,
101    P: Parser,
102    I: Iterator<Item = Option<&'a str>>,
103    F: FnOnce() -> Option<NullBuffer>,
104>(
105    iter: I,
106    cast_options: &CastOptions,
107    nulls: F,
108) -> Result<ArrayRef, ArrowError> {
109    let array = if cast_options.safe {
110        let iter = iter.map(|x| x.and_then(P::parse));
111
112        // Benefit:
113        //     20% performance improvement
114        // Soundness:
115        //     The iterator is trustedLen because it comes from an `StringArray`.
116        unsafe { PrimitiveArray::<P>::from_trusted_len_iter(iter) }
117    } else {
118        let v = iter
119            .map(|x| match x {
120                Some(v) => P::parse(v).ok_or_else(|| {
121                    ArrowError::CastError(format!(
122                        "Cannot cast string '{v}' to value of {} type",
123                        P::DATA_TYPE
124                    ))
125                }),
126                None => Ok(P::Native::default()),
127            })
128            .collect::<Result<Vec<_>, ArrowError>>()?;
129        PrimitiveArray::try_new(v.into(), nulls())?
130    };
131
132    Ok(Arc::new(array) as ArrayRef)
133}
134
135/// Casts generic string arrays to an ArrowTimestampType (TimeStampNanosecondArray, etc.)
136pub(crate) fn cast_string_to_timestamp<O: OffsetSizeTrait, T: ArrowTimestampType>(
137    array: &dyn Array,
138    to_tz: Option<&Arc<str>>,
139    cast_options: &CastOptions,
140) -> Result<ArrayRef, ArrowError> {
141    let array = array.as_string::<O>();
142    let out: PrimitiveArray<T> = match to_tz {
143        Some(tz) => {
144            let tz: Tz = tz.as_ref().parse()?;
145            cast_string_to_timestamp_impl(array.iter(), &tz, cast_options)?
146        }
147        None => cast_string_to_timestamp_impl(array.iter(), &Utc, cast_options)?,
148    };
149    Ok(Arc::new(out.with_timezone_opt(to_tz.cloned())))
150}
151
152/// Casts string view arrays to an ArrowTimestampType (TimeStampNanosecondArray, etc.)
153pub(crate) fn cast_view_to_timestamp<T: ArrowTimestampType>(
154    array: &dyn Array,
155    to_tz: Option<&Arc<str>>,
156    cast_options: &CastOptions,
157) -> Result<ArrayRef, ArrowError> {
158    let array = array.as_string_view();
159    let out: PrimitiveArray<T> = match to_tz {
160        Some(tz) => {
161            let tz: Tz = tz.as_ref().parse()?;
162            cast_string_to_timestamp_impl(array.iter(), &tz, cast_options)?
163        }
164        None => cast_string_to_timestamp_impl(array.iter(), &Utc, cast_options)?,
165    };
166    Ok(Arc::new(out.with_timezone_opt(to_tz.cloned())))
167}
168
169fn cast_string_to_timestamp_impl<
170    'a,
171    I: Iterator<Item = Option<&'a str>>,
172    T: ArrowTimestampType,
173    Tz: TimeZone,
174>(
175    iter: I,
176    tz: &Tz,
177    cast_options: &CastOptions,
178) -> Result<PrimitiveArray<T>, ArrowError> {
179    if cast_options.safe {
180        let iter = iter.map(|v| {
181            let naive = string_to_datetime(tz, v?).ok()?.naive_utc();
182            T::from_naive_datetime(naive, None)
183        });
184        // Benefit:
185        //     20% performance improvement
186        // Soundness:
187        //     The iterator is trustedLen because it comes from an `StringArray`.
188
189        Ok(unsafe { PrimitiveArray::from_trusted_len_iter(iter) })
190    } else {
191        let vec = iter
192            .map(|v| {
193                v.map(|v| {
194                    let naive = string_to_datetime(tz, v)?.naive_utc();
195                    T::from_naive_datetime(naive, None).ok_or_else(|| match T::UNIT {
196                        TimeUnit::Nanosecond => ArrowError::CastError(format!(
197                            "Overflow converting {naive} to Nanosecond. The dates that can be represented as nanoseconds have to be between 1677-09-21T00:12:44.0 and 2262-04-11T23:47:16.854775804"
198                        )),
199                        _ => ArrowError::CastError(format!(
200                            "Overflow converting {naive} to {:?}",
201                            T::UNIT
202                        ))
203                    })
204                })
205                    .transpose()
206            })
207            .collect::<Result<Vec<Option<i64>>, _>>()?;
208
209        // Benefit:
210        //     20% performance improvement
211        // Soundness:
212        //     The iterator is trustedLen because it comes from an `StringArray`.
213        Ok(unsafe { PrimitiveArray::from_trusted_len_iter(vec.iter()) })
214    }
215}
216
217pub(crate) fn cast_string_to_interval<Offset, F, ArrowType>(
218    array: &dyn Array,
219    cast_options: &CastOptions,
220    parse_function: F,
221) -> Result<ArrayRef, ArrowError>
222where
223    Offset: OffsetSizeTrait,
224    ArrowType: ArrowPrimitiveType,
225    F: Fn(&str) -> Result<ArrowType::Native, ArrowError> + Copy,
226{
227    let string_array = array
228        .as_any()
229        .downcast_ref::<GenericStringArray<Offset>>()
230        .unwrap();
231    cast_string_to_interval_impl::<_, ArrowType, F>(
232        string_array.iter(),
233        cast_options,
234        parse_function,
235    )
236}
237
238pub(crate) fn cast_string_to_year_month_interval<Offset: OffsetSizeTrait>(
239    array: &dyn Array,
240    cast_options: &CastOptions,
241) -> Result<ArrayRef, ArrowError> {
242    cast_string_to_interval::<Offset, _, IntervalYearMonthType>(
243        array,
244        cast_options,
245        parse_interval_year_month,
246    )
247}
248
249pub(crate) fn cast_string_to_day_time_interval<Offset: OffsetSizeTrait>(
250    array: &dyn Array,
251    cast_options: &CastOptions,
252) -> Result<ArrayRef, ArrowError> {
253    cast_string_to_interval::<Offset, _, IntervalDayTimeType>(
254        array,
255        cast_options,
256        parse_interval_day_time,
257    )
258}
259
260pub(crate) fn cast_string_to_month_day_nano_interval<Offset: OffsetSizeTrait>(
261    array: &dyn Array,
262    cast_options: &CastOptions,
263) -> Result<ArrayRef, ArrowError> {
264    cast_string_to_interval::<Offset, _, IntervalMonthDayNanoType>(
265        array,
266        cast_options,
267        parse_interval_month_day_nano,
268    )
269}
270
271pub(crate) fn cast_view_to_interval<F, ArrowType>(
272    array: &dyn Array,
273    cast_options: &CastOptions,
274    parse_function: F,
275) -> Result<ArrayRef, ArrowError>
276where
277    ArrowType: ArrowPrimitiveType,
278    F: Fn(&str) -> Result<ArrowType::Native, ArrowError> + Copy,
279{
280    let string_view_array = array.as_any().downcast_ref::<StringViewArray>().unwrap();
281    cast_string_to_interval_impl::<_, ArrowType, F>(
282        string_view_array.iter(),
283        cast_options,
284        parse_function,
285    )
286}
287
288pub(crate) fn cast_view_to_year_month_interval(
289    array: &dyn Array,
290    cast_options: &CastOptions,
291) -> Result<ArrayRef, ArrowError> {
292    cast_view_to_interval::<_, IntervalYearMonthType>(
293        array,
294        cast_options,
295        parse_interval_year_month,
296    )
297}
298
299pub(crate) fn cast_view_to_day_time_interval(
300    array: &dyn Array,
301    cast_options: &CastOptions,
302) -> Result<ArrayRef, ArrowError> {
303    cast_view_to_interval::<_, IntervalDayTimeType>(array, cast_options, parse_interval_day_time)
304}
305
306pub(crate) fn cast_view_to_month_day_nano_interval(
307    array: &dyn Array,
308    cast_options: &CastOptions,
309) -> Result<ArrayRef, ArrowError> {
310    cast_view_to_interval::<_, IntervalMonthDayNanoType>(
311        array,
312        cast_options,
313        parse_interval_month_day_nano,
314    )
315}
316
317fn cast_string_to_interval_impl<'a, I, ArrowType, F>(
318    iter: I,
319    cast_options: &CastOptions,
320    parse_function: F,
321) -> Result<ArrayRef, ArrowError>
322where
323    I: Iterator<Item = Option<&'a str>>,
324    ArrowType: ArrowPrimitiveType,
325    F: Fn(&str) -> Result<ArrowType::Native, ArrowError> + Copy,
326{
327    let interval_array = if cast_options.safe {
328        let iter = iter.map(|v| parse_function(v?).ok());
329
330        // Benefit:
331        //     20% performance improvement
332        // Soundness:
333        //     The iterator is trustedLen because it comes from an `StringArray`.
334        unsafe { PrimitiveArray::<ArrowType>::from_trusted_len_iter(iter) }
335    } else {
336        let vec = iter
337            .map(|v| v.map(parse_function).transpose())
338            .collect::<Result<Vec<_>, ArrowError>>()?;
339
340        // Benefit:
341        //     20% performance improvement
342        // Soundness:
343        //     The iterator is trustedLen because it comes from an `StringArray`.
344        unsafe { PrimitiveArray::<ArrowType>::from_trusted_len_iter(vec) }
345    };
346    Ok(Arc::new(interval_array) as ArrayRef)
347}
348
349/// A specified helper to cast from `GenericBinaryArray` to `GenericStringArray` when they have same
350/// offset size so re-encoding offset is unnecessary.
351fn extend_valid_utf8<'a, B, I>(builder: &mut B, iter: I)
352where
353    B: Extend<Option<&'a str>>,
354    I: Iterator<Item = Option<&'a [u8]>>,
355{
356    builder.extend(iter.map(|value| std::str::from_utf8(value?).ok()));
357}
358
359pub(crate) fn cast_binary_to_string<O: OffsetSizeTrait>(
360    array: &dyn Array,
361    cast_options: &CastOptions,
362) -> Result<ArrayRef, ArrowError> {
363    let array = array
364        .as_any()
365        .downcast_ref::<GenericByteArray<GenericBinaryType<O>>>()
366        .unwrap();
367
368    match GenericStringArray::<O>::try_from_binary(array.clone()) {
369        Ok(a) => Ok(Arc::new(a)),
370        Err(e) => match cast_options.safe {
371            true => {
372                // Fallback to slow method to convert invalid sequences to nulls.
373                // Reserve only the bytes the offsets span, not the whole values buffer.
374                let offsets = array.offsets();
375                let data_len = (offsets.last() - offsets.first()).as_usize();
376                let mut builder = GenericStringBuilder::<O>::with_capacity(array.len(), data_len);
377
378                extend_valid_utf8(&mut builder, array.iter());
379                Ok(Arc::new(builder.finish()))
380            }
381            false => Err(e),
382        },
383    }
384}
385
386pub(crate) fn cast_binary_view_to_string_view(
387    array: &dyn Array,
388    cast_options: &CastOptions,
389) -> Result<ArrayRef, ArrowError> {
390    let array = array.as_binary_view();
391
392    match array.clone().to_string_view() {
393        Ok(result) => Ok(Arc::new(result)),
394        Err(error) => match cast_options.safe {
395            true => {
396                let mut builder = StringViewBuilder::with_capacity(array.len());
397                extend_valid_utf8(&mut builder, array.iter());
398                Ok(Arc::new(builder.finish()))
399            }
400            false => Err(error),
401        },
402    }
403}
404
405/// Casts string to boolean
406fn cast_string_to_boolean<'a, StrArray>(
407    array: &StrArray,
408    cast_options: &CastOptions,
409) -> Result<ArrayRef, ArrowError>
410where
411    StrArray: StringArrayType<'a>,
412{
413    let output_array = array
414        .iter()
415        .map(|value| match value {
416            Some(value) => cast_single_string_to_boolean(value, cast_options),
417            None => Ok(None),
418        })
419        .collect::<Result<BooleanArray, _>>()?;
420
421    Ok(Arc::new(output_array))
422}
423
424#[inline]
425fn cast_single_string_to_boolean(
426    value: &str,
427    cast_options: &CastOptions,
428) -> Result<Option<bool>, ArrowError> {
429    match value.to_ascii_lowercase().trim() {
430        "t" | "tr" | "tru" | "true" | "y" | "ye" | "yes" | "on" | "1" => Ok(Some(true)),
431        "f" | "fa" | "fal" | "fals" | "false" | "n" | "no" | "of" | "off" | "0" => Ok(Some(false)),
432        invalid_value => match cast_options.safe {
433            true => Ok(None),
434            false => Err(ArrowError::CastError(format!(
435                "Cannot cast value '{invalid_value}' to value of Boolean type",
436            ))),
437        },
438    }
439}
440
441/// Cast a single string to boolean with default cast option(safe=true).
442pub fn cast_single_string_to_boolean_default(value: &str) -> Option<bool> {
443    cast_single_string_to_boolean(value, &CastOptions::default())
444        .ok()
445        .flatten()
446}
447
448pub(crate) fn cast_utf8_to_boolean<OffsetSize>(
449    from: &dyn Array,
450    cast_options: &CastOptions,
451) -> Result<ArrayRef, ArrowError>
452where
453    OffsetSize: OffsetSizeTrait,
454{
455    let array = from
456        .as_any()
457        .downcast_ref::<GenericStringArray<OffsetSize>>()
458        .unwrap();
459
460    cast_string_to_boolean(&array, cast_options)
461}
462
463pub(crate) fn cast_utf8view_to_boolean(
464    from: &dyn Array,
465    cast_options: &CastOptions,
466) -> Result<ArrayRef, ArrowError> {
467    let array = from.as_any().downcast_ref::<StringViewArray>().unwrap();
468
469    cast_string_to_boolean(&array, cast_options)
470}
arrow_cast/cast/string.rs

arrow_cast/cast/
string.rs
