lance_arrow/
cast.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::sync::Arc;
5
6use arrow_array::{Array, ArrayRef, FixedSizeListArray};
7use arrow_cast::CastOptions;
8use arrow_schema::{ArrowError, DataType};
9
10/// Customized [`arrow_cast::can_cast_types`] that handles cases not supported upstream yet.
11pub fn can_cast_types(from_type: &DataType, to_type: &DataType) -> bool {
12    use DataType::*;
13    match (from_type, to_type) {
14        // TODO: remove this once Arrow supports this upstream.
15        // https://github.com/apache/arrow-rs/pull/5360
16        (FixedSizeList(from_field, size_from), FixedSizeList(to_field, size_to)) => {
17            size_from == size_to && can_cast_types(from_field.data_type(), to_field.data_type())
18        }
19        // TODO: support bfloat16 cast?
20        _ => arrow_cast::can_cast_types(from_type, to_type),
21    }
22}
23
24/// Customized [`arrow_cast::cast_with_options`] that handles cases not supported upstream yet.
25pub fn cast_with_options(
26    array: &dyn Array,
27    to_type: &DataType,
28    cast_options: &CastOptions,
29) -> Result<ArrayRef, ArrowError> {
30    use DataType::*;
31    match (array.data_type(), to_type) {
32        (FixedSizeList(_, size_from), FixedSizeList(to_field, size_to)) if size_from == size_to => {
33            let array = array.as_any().downcast_ref::<FixedSizeListArray>().unwrap();
34            let values = cast_with_options(array.values(), to_field.data_type(), cast_options)?;
35            Ok(Arc::new(FixedSizeListArray::try_new(
36                to_field.clone(),
37                *size_from,
38                values,
39                array.nulls().cloned(),
40            )?))
41        }
42        _ => arrow_cast::cast_with_options(array, to_type, cast_options),
43    }
44}
lance_arrow/cast.rs

lance_arrow/
cast.rs
