arrow_cast/cast/
run_array.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
5use crate::cast::*;
6use arrow_ord::partition::partition;
7
8/// Attempts to cast a `RunArray` with index type K into
9/// `to_type` for supported types.
10pub(crate) fn run_end_encoded_cast<K: RunEndIndexType>(
11    array: &dyn Array,
12    to_type: &DataType,
13    cast_options: &CastOptions,
14) -> Result<ArrayRef, ArrowError> {
15    match array.data_type() {
16        DataType::RunEndEncoded(_, _) => {
17            let run_array = array.as_run::<K>();
18
19            match to_type {
20                // Stay as RunEndEncoded, cast only the values
21                DataType::RunEndEncoded(target_index_field, target_value_field) => {
22                    let values = run_array.values_slice();
23                    let cast_values = cast_with_options(
24                        values.as_ref(),
25                        target_value_field.data_type(),
26                        cast_options,
27                    )?;
28
29                    let run_ends_array =
30                        PrimitiveArray::<K>::from_iter_values(run_array.run_ends().sliced_values());
31                    let cast_run_ends = cast_with_options(
32                        &run_ends_array,
33                        target_index_field.data_type(),
34                        cast_options,
35                    )?;
36                    let new_run_array: ArrayRef = match target_index_field.data_type() {
37                        DataType::Int16 => {
38                            let re = cast_run_ends.as_primitive::<Int16Type>();
39                            Arc::new(RunArray::<Int16Type>::try_new(re, cast_values.as_ref())?)
40                        }
41                        DataType::Int32 => {
42                            let re = cast_run_ends.as_primitive::<Int32Type>();
43                            Arc::new(RunArray::<Int32Type>::try_new(re, cast_values.as_ref())?)
44                        }
45                        DataType::Int64 => {
46                            let re = cast_run_ends.as_primitive::<Int64Type>();
47                            Arc::new(RunArray::<Int64Type>::try_new(re, cast_values.as_ref())?)
48                        }
49                        _ => {
50                            return Err(ArrowError::InvalidArgumentError(
51                                "Run-end type must be i16, i32, or i64".to_string(),
52                            ));
53                        }
54                    };
55                    Ok(new_run_array)
56                }
57
58                // Expand to logical form
59                _ => {
60                    let values = run_array.values();
61                    let len = run_array.len();
62                    let offset = run_array.offset();
63                    let run_ends = run_array.run_ends().values();
64
65                    let mut indices = Vec::with_capacity(len);
66                    let mut physical_idx = run_array.get_start_physical_index();
67
68                    for logical_idx in offset..offset + len {
69                        if logical_idx == run_ends[physical_idx].as_usize() {
70                            // If the logical index is equal to the (next) run end, increment the physical index,
71                            // since we are at the end of a run.
72                            physical_idx += 1;
73                        }
74                        indices.push(physical_idx as i32);
75                    }
76
77                    let taken = take(&values, &Int32Array::from_iter_values(indices), None)?;
78                    if taken.data_type() != to_type {
79                        cast_with_options(taken.as_ref(), to_type, cast_options)
80                    } else {
81                        Ok(taken)
82                    }
83                }
84            }
85        }
86
87        _ => unreachable!(
88            "cast_with_options only calls run_end_encoded_cast for RunEndEncoded source arrays"
89        ),
90    }
91}
92
93/// Attempts to encode an array into a `RunArray` with index type K
94/// and value type `value_type`
95pub(crate) fn cast_to_run_end_encoded<K: RunEndIndexType>(
96    array: &ArrayRef,
97    value_type: &DataType,
98    cast_options: &CastOptions,
99) -> Result<ArrayRef, ArrowError> {
100    let mut run_ends_builder = PrimitiveBuilder::<K>::new();
101
102    // Cast the input array to the target value type if necessary
103    let cast_array = if array.data_type() == value_type {
104        array
105    } else {
106        &cast_with_options(array, value_type, cast_options)?
107    };
108
109    // Return early if the array to cast is empty
110    if cast_array.is_empty() {
111        let empty_run_ends = run_ends_builder.finish();
112        let empty_values = make_array(ArrayData::new_empty(value_type));
113        return Ok(Arc::new(RunArray::<K>::try_new(
114            &empty_run_ends,
115            empty_values.as_ref(),
116        )?));
117    }
118
119    // REE arrays are handled by run_end_encoded_cast
120    if let DataType::RunEndEncoded(_, _) = array.data_type() {
121        unreachable!(
122            "cast_with_options handles RunEndEncoded source arrays with run_end_encoded_cast first"
123        );
124    }
125
126    // Partition the array to identify runs of consecutive equal values
127    let partitions = partition(&[Arc::clone(cast_array)])?;
128    let size = partitions.len();
129    let mut run_ends = Vec::with_capacity(size);
130    let mut values_indexes = Vec::with_capacity(size);
131    let mut last_partition_end = 0;
132    for partition in partitions.ranges() {
133        values_indexes.push(last_partition_end);
134        run_ends.push(partition.end);
135        last_partition_end = partition.end;
136    }
137
138    // Build the run_ends array
139    for run_end in run_ends {
140        run_ends_builder.append_value(K::Native::from_usize(run_end).ok_or_else(|| {
141            ArrowError::CastError(format!("Run end index out of range: {run_end}"))
142        })?);
143    }
144    let run_ends_array = run_ends_builder.finish();
145    // Build the values array by taking elements at the run start positions
146    let indices = PrimitiveArray::<UInt32Type>::from_iter_values(
147        values_indexes.iter().map(|&idx| idx as u32),
148    );
149    let values_array = take(&cast_array, &indices, None)?;
150
151    // Create and return the RunArray
152    let run_array = RunArray::<K>::try_new(&run_ends_array, values_array.as_ref())?;
153    Ok(Arc::new(run_array))
154}
arrow_cast/cast/run_array.rs

arrow_cast/cast/
run_array.rs
