vortex_array/arrays/map/compute/
cast.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::sync::Arc;
5
6use vortex_error::VortexResult;
7use vortex_error::vortex_bail;
8
9use crate::ArrayRef;
10use crate::array::ArrayView;
11use crate::arrays::ListView;
12use crate::arrays::map::Map;
13use crate::arrays::map::MapArrayExt;
14use crate::arrays::map::MapArraySlotsExt;
15use crate::arrays::map::compute::rebuild_map_from_array;
16use crate::dtype::DType;
17use crate::dtype::MapDType;
18use crate::executor::ExecutionCtx;
19use crate::scalar_fn::fns::cast::CastKernel;
20use crate::scalar_fn::fns::cast::CastReduce;
21
22fn prepare_map_cast_target(
23    array: ArrayView<'_, Map>,
24    dtype: &DType,
25) -> VortexResult<Option<(MapDType, DType)>> {
26    let Some(target_map_dtype) = dtype.as_map_opt() else {
27        return Ok(None);
28    };
29
30    if target_map_dtype.keys_sorted() && !array.keys_sorted() {
31        vortex_bail!(
32            "Cannot cast {} to {dtype}: source does not assert sorted map keys",
33            array.dtype()
34        );
35    }
36
37    let target_entries_dtype = DType::List(
38        Arc::new(target_map_dtype.entries_dtype()),
39        dtype.nullability(),
40    );
41    Ok(Some((target_map_dtype.clone(), target_entries_dtype)))
42}
43
44impl CastReduce for Map {
45    fn cast(array: ArrayView<'_, Self>, dtype: &DType) -> VortexResult<Option<ArrayRef>> {
46        let Some((target_map_dtype, target_entries_dtype)) = prepare_map_cast_target(array, dtype)?
47        else {
48            return Ok(None);
49        };
50
51        let Some(entries) = <ListView as CastReduce>::cast(
52            array.entries().as_::<ListView>(),
53            &target_entries_dtype,
54        )?
55        else {
56            return Ok(None);
57        };
58
59        rebuild_map_from_array(target_map_dtype, entries).map(Some)
60    }
61}
62
63impl CastKernel for Map {
64    fn cast(
65        array: ArrayView<'_, Self>,
66        dtype: &DType,
67        ctx: &mut ExecutionCtx,
68    ) -> VortexResult<Option<ArrayRef>> {
69        let Some((target_map_dtype, target_entries_dtype)) = prepare_map_cast_target(array, dtype)?
70        else {
71            return Ok(None);
72        };
73
74        let Some(entries) = <ListView as CastKernel>::cast(
75            array.entries().as_::<ListView>(),
76            &target_entries_dtype,
77            ctx,
78        )?
79        else {
80            return Ok(None);
81        };
82
83        rebuild_map_from_array(target_map_dtype, entries).map(Some)
84    }
85}
vortex_array/arrays/map/compute/cast.rs

vortex_array/arrays/map/compute/
cast.rs
