Skip to main content

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