arrow_cast/cast/
string.rs

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

arrow_cast/cast/
string.rs
