Skip to main content

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