Skip to main content

arrow_cast/cast/
binary.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
5//! Cast support for binary, fixed-size binary and byte-view arrays.
6
7use arrow_array::{builder::*, cast::*, types::*, *};
8use arrow_buffer::{Buffer, OffsetBuffer};
9use arrow_data::{ArrayData, ByteView};
10use arrow_schema::{ArrowError, DataType};
11use num_traits::{NumCast, ToPrimitive};
12use std::sync::Arc;
13
14use super::CastOptions;
15
16pub(crate) fn cast_numeric_to_binary<FROM: ArrowPrimitiveType, O: OffsetSizeTrait>(
17    array: &dyn Array,
18) -> Result<ArrayRef, ArrowError> {
19    let array = array.as_primitive::<FROM>();
20    let size = std::mem::size_of::<FROM::Native>();
21    let offsets = OffsetBuffer::from_repeated_length(size, array.len());
22    Ok(Arc::new(GenericBinaryArray::<O>::try_new(
23        offsets,
24        array.values().inner().clone(),
25        array.nulls().cloned(),
26    )?))
27}
28
29/// Helper function to cast from one `BinaryArray` or 'LargeBinaryArray' to 'FixedSizeBinaryArray'.
30pub(crate) fn cast_binary_to_fixed_size_binary<O: OffsetSizeTrait>(
31    array: &dyn Array,
32    byte_width: i32,
33    cast_options: &CastOptions,
34) -> Result<ArrayRef, ArrowError> {
35    let array = array.as_binary::<O>();
36    let mut builder = FixedSizeBinaryBuilder::with_capacity(array.len(), byte_width);
37
38    for i in 0..array.len() {
39        if array.is_null(i) {
40            builder.append_null();
41        } else {
42            match builder.append_value(array.value(i)) {
43                Ok(()) => {}
44                Err(e) => match cast_options.safe {
45                    true => builder.append_null(),
46                    false => return Err(e),
47                },
48            }
49        }
50    }
51
52    Ok(Arc::new(builder.finish()))
53}
54
55/// Helper function to cast from 'FixedSizeBinaryArray' to one `BinaryArray` or 'LargeBinaryArray'.
56/// If the target one is too large for the source array it will return an Error.
57pub(crate) fn cast_fixed_size_binary_to_binary<O: OffsetSizeTrait>(
58    array: &dyn Array,
59    byte_width: i32,
60) -> Result<ArrayRef, ArrowError> {
61    let array = array
62        .as_any()
63        .downcast_ref::<FixedSizeBinaryArray>()
64        .unwrap();
65
66    let offsets: i128 = byte_width as i128 * array.len() as i128;
67
68    let is_binary = matches!(GenericBinaryType::<O>::DATA_TYPE, DataType::Binary);
69    if is_binary && offsets > i32::MAX as i128 {
70        return Err(ArrowError::ComputeError(
71            "FixedSizeBinary array too large to cast to Binary array".to_string(),
72        ));
73    } else if !is_binary && offsets > i64::MAX as i128 {
74        return Err(ArrowError::ComputeError(
75            "FixedSizeBinary array too large to cast to LargeBinary array".to_string(),
76        ));
77    }
78
79    let mut builder = GenericBinaryBuilder::<O>::with_capacity(array.len(), array.len());
80
81    for i in 0..array.len() {
82        if array.is_null(i) {
83            builder.append_null();
84        } else {
85            builder.append_value(array.value(i));
86        }
87    }
88
89    Ok(Arc::new(builder.finish()))
90}
91
92pub(crate) fn cast_fixed_size_binary_to_binary_view(
93    array: &dyn Array,
94    _byte_width: i32,
95) -> Result<ArrayRef, ArrowError> {
96    let array = array
97        .as_any()
98        .downcast_ref::<FixedSizeBinaryArray>()
99        .unwrap();
100
101    let mut builder = BinaryViewBuilder::with_capacity(array.len());
102    for i in 0..array.len() {
103        if array.is_null(i) {
104            builder.append_null();
105        } else {
106            builder.append_value(array.value(i));
107        }
108    }
109
110    Ok(Arc::new(builder.finish()))
111}
112
113/// Helper function to cast from one `ByteArrayType` to another and vice versa.
114/// If the target one (e.g., `LargeUtf8`) is too large for the source array it will return an Error.
115pub(crate) fn cast_byte_container<FROM, TO>(array: &dyn Array) -> Result<ArrayRef, ArrowError>
116where
117    FROM: ByteArrayType,
118    TO: ByteArrayType<Native = FROM::Native>,
119    FROM::Offset: OffsetSizeTrait + ToPrimitive,
120    TO::Offset: OffsetSizeTrait + NumCast,
121{
122    let data = array.to_data();
123    assert_eq!(data.data_type(), &FROM::DATA_TYPE);
124    let str_values_buf = data.buffers()[1].clone();
125    let offsets = data.buffers()[0].typed_data::<FROM::Offset>();
126
127    let mut cast_offsets = Vec::<TO::Offset>::with_capacity(offsets.len());
128    offsets
129        .iter()
130        .try_for_each::<_, Result<_, ArrowError>>(|offset| {
131            let offset =
132                <<TO as ByteArrayType>::Offset as NumCast>::from(*offset).ok_or_else(|| {
133                    ArrowError::ComputeError(format!(
134                        "{}{} array too large to cast to {}{} array",
135                        FROM::Offset::PREFIX,
136                        FROM::PREFIX,
137                        TO::Offset::PREFIX,
138                        TO::PREFIX
139                    ))
140                })?;
141            cast_offsets.push(offset);
142            Ok(())
143        })?;
144
145    let offset_buffer = Buffer::from_vec(cast_offsets);
146
147    let dtype = TO::DATA_TYPE;
148
149    let builder = ArrayData::builder(dtype)
150        .offset(array.offset())
151        .len(array.len())
152        .add_buffer(offset_buffer)
153        .add_buffer(str_values_buf)
154        .nulls(data.nulls().cloned());
155
156    let array_data = unsafe { builder.build_unchecked() };
157
158    Ok(Arc::new(GenericByteArray::<TO>::from(array_data)))
159}
160
161/// Helper function to cast from one `ByteViewType` array to `ByteArrayType` array.
162pub(crate) fn cast_view_to_byte<FROM, TO>(array: &dyn Array) -> Result<ArrayRef, ArrowError>
163where
164    FROM: ByteViewType,
165    TO: ByteArrayType,
166    FROM::Native: AsRef<TO::Native>,
167{
168    let data = array.to_data();
169    let view_array = GenericByteViewArray::<FROM>::from(data);
170
171    let len = view_array.len();
172    let bytes = view_array
173        .views()
174        .iter()
175        .map(|v| ByteView::from(*v).length as usize)
176        .sum::<usize>();
177
178    let mut byte_array_builder = GenericByteBuilder::<TO>::with_capacity(len, bytes);
179
180    for val in &view_array {
181        byte_array_builder.append_option(val);
182    }
183
184    Ok(Arc::new(byte_array_builder.finish()))
185}