vortex_arrow/
run_end_import.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use arrow_array::RunArray;
5use arrow_array::types::RunEndIndexType;
6use vortex_array::ArrayRef;
7use vortex_array::IntoArray;
8use vortex_array::VortexSessionExecute;
9use vortex_array::arrays::PrimitiveArray;
10use vortex_array::arrays::primitive::PrimitiveArrayExt;
11use vortex_array::dtype::NativePType;
12use vortex_array::legacy_session;
13use vortex_array::validity::Validity;
14use vortex_buffer::Buffer;
15use vortex_error::VortexResult;
16use vortex_runend::RunEndData;
17use vortex_runend::ops::find_physical_index;
18use vortex_runend::ops::find_slice_end_index;
19
20use crate::FromArrowArray;
21
22impl<R: RunEndIndexType> FromArrowArray<&RunArray<R>> for RunEndData
23where
24    R::Native: NativePType,
25{
26    #[allow(clippy::disallowed_methods)]
27    fn from_arrow(array: &RunArray<R>, nullable: bool) -> VortexResult<Self> {
28        let offset = array.run_ends().offset();
29        let len = array.run_ends().len();
30        let ends_buf =
31            Buffer::<R::Native>::from_arrow_scalar_buffer(array.run_ends().inner().clone());
32        let ends = PrimitiveArray::new(ends_buf, Validity::NonNullable)
33            .reinterpret_cast(R::Native::PTYPE.to_unsigned());
34        let values = ArrayRef::from_arrow(array.values().as_ref(), nullable)?;
35
36        let ends_array = PrimitiveArray::from_buffer_handle(
37            ends.buffer_handle().clone(),
38            ends.ptype(),
39            ends.validity()?,
40        )
41        .into_array();
42        let mut ctx = legacy_session().create_execution_ctx();
43
44        let (ends_slice, values_slice) = if offset == 0 && len == array.run_ends().max_value() {
45            (ends_array, values)
46        } else {
47            let slice_begin = find_physical_index(&ends_array, offset, &mut ctx)?;
48            let slice_end = find_slice_end_index(&ends_array, offset + len, &mut ctx)?;
49
50            (
51                ends_array.slice(slice_begin..slice_end)?,
52                values.slice(slice_begin..slice_end)?,
53            )
54        };
55
56        // SAFETY: arrow-rs enforces the RunEndArray invariants, we inherit their guarantees.
57        RunEndData::validate_parts(&ends_slice, &values_slice, offset, len, &mut ctx)?;
58        Ok(unsafe { RunEndData::new_unchecked(offset) })
59    }
60}
61
62#[cfg(test)]
63mod tests {
64    use std::sync::Arc;
65    use std::sync::LazyLock;
66
67    use arrow_array::Array as _;
68    use arrow_array::Float64Array;
69    use arrow_array::Int32Array;
70    use arrow_array::Int64Array;
71    use arrow_array::RunArray;
72    use arrow_array::types::Int32Type;
73    use arrow_array::types::Int64Type;
74    use arrow_array::types::RunEndIndexType;
75    use arrow_schema::DataType;
76    use arrow_schema::Field;
77    use rstest::rstest;
78    use vortex_array::ArrayRef;
79    use vortex_array::IntoArray as _;
80    use vortex_array::VortexSessionExecute as _;
81    use vortex_array::arrays::PrimitiveArray;
82    use vortex_array::arrays::primitive::PrimitiveArrayExt;
83    use vortex_array::assert_arrays_eq;
84    use vortex_array::dtype::DType;
85    use vortex_array::dtype::NativePType;
86    use vortex_array::dtype::Nullability;
87    use vortex_array::dtype::PType;
88    use vortex_array::validity::Validity;
89    use vortex_buffer::Buffer;
90    use vortex_buffer::buffer;
91    use vortex_error::VortexResult;
92    use vortex_runend::RunEnd;
93    use vortex_runend::RunEndArray;
94    use vortex_runend::ops::find_physical_index;
95    use vortex_runend::ops::find_slice_end_index;
96    use vortex_session::VortexSession;
97
98    use crate::ArrowSessionExt;
99    use crate::FromArrowArray;
100
101    static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
102        let session = vortex_array::array_session();
103        vortex_runend::initialize(&session);
104        session
105    });
106
107    fn decode_run_array<R: RunEndIndexType>(
108        array: &RunArray<R>,
109        nullable: bool,
110    ) -> VortexResult<RunEndArray>
111    where
112        R::Native: NativePType,
113    {
114        let offset = array.run_ends().offset();
115        let len = array.run_ends().len();
116        let ends_buf =
117            Buffer::<R::Native>::from_arrow_scalar_buffer(array.run_ends().inner().clone());
118        let ends = PrimitiveArray::new(ends_buf, Validity::NonNullable)
119            .reinterpret_cast(R::Native::PTYPE.to_unsigned());
120        let values = ArrayRef::from_arrow(array.values().as_ref(), nullable)?;
121
122        let ends_array = PrimitiveArray::from_buffer_handle(
123            ends.buffer_handle().clone(),
124            ends.ptype(),
125            ends.validity()?,
126        )
127        .into_array();
128        let mut ctx = SESSION.create_execution_ctx();
129        let (ends_slice, values_slice) = if offset == 0 && len == array.run_ends().max_value() {
130            (ends_array, values)
131        } else {
132            let slice_begin = find_physical_index(&ends_array, offset, &mut ctx)?;
133            let slice_end = find_slice_end_index(&ends_array, offset + len, &mut ctx)?;
134
135            (
136                ends_array.slice(slice_begin..slice_end)?,
137                values.slice(slice_begin..slice_end)?,
138            )
139        };
140
141        RunEnd::try_new_offset_length(ends_slice, values_slice, offset, array.len(), &mut ctx)
142    }
143
144    #[test]
145    fn test_arrow_run_array_to_vortex() -> VortexResult<()> {
146        // Create an Arrow RunArray with UInt32 run ends and Int32 values
147        // Run ends: [3, 5, 8] means runs of length 3, 2, 3
148        // Values: [10, 20, 30] means values 10, 10, 10, 20, 20, 30, 30, 30
149        let run_ends = Int32Array::from(vec![3i32, 5, 8]);
150        let values = Int32Array::from(vec![10, 20, 30]);
151        let arrow_run_array = RunArray::<Int32Type>::try_new(&run_ends, &values)?;
152
153        // Convert to Vortex
154        let vortex_array = decode_run_array(&arrow_run_array, false)?;
155
156        assert_arrays_eq!(
157            vortex_array.into_array(),
158            buffer![10i32, 10, 10, 20, 20, 30, 30, 30].into_array(),
159            &mut SESSION.create_execution_ctx()
160        );
161        Ok(())
162    }
163
164    #[test]
165    fn test_arrow_run_array_with_nulls_to_vortex() -> VortexResult<()> {
166        // Create an Arrow RunArray with nullable values
167        let run_ends = Int32Array::from(vec![2i32, 4, 6]);
168        let values = Int32Array::from(vec![Some(100), None, Some(300)]);
169        let arrow_run_array = RunArray::<Int32Type>::try_new(&run_ends, &values)?;
170
171        // Convert to Vortex with nullable=true
172        let vortex_array = decode_run_array(&arrow_run_array, true)?;
173
174        assert_arrays_eq!(
175            vortex_array.into_array(),
176            PrimitiveArray::from_option_iter([
177                Some(100i32),
178                Some(100i32),
179                None,
180                None,
181                Some(300i32),
182                Some(300i32)
183            ]),
184            &mut SESSION.create_execution_ctx()
185        );
186        Ok(())
187    }
188
189    #[test]
190    fn test_arrow_run_array_with_different_types() -> VortexResult<()> {
191        // Test with UInt64 run ends and Float64 values
192        let run_ends = Int64Array::from(vec![1i64, 3, 4]);
193        let values = Float64Array::from(vec![1.5, 2.5, 3.5]);
194        let arrow_run_array = RunArray::<Int64Type>::try_new(&run_ends, &values)?;
195
196        // Convert to Vortex
197        let vortex_array = decode_run_array(&arrow_run_array, false)?;
198
199        assert_arrays_eq!(
200            vortex_array,
201            buffer![1.5f64, 2.5, 2.5, 3.5].into_array(),
202            &mut SESSION.create_execution_ctx()
203        );
204        Ok(())
205    }
206
207    #[test]
208    fn test_sliced_arrow_run_array_to_vortex() -> VortexResult<()> {
209        // Create an Arrow RunArray with run ends and values
210        // Run ends: [2, 5, 8, 10] means runs of length 2, 3, 3, 2
211        // Values: [100, 200, 300, 400] means: 100, 100, 200, 200, 200, 300, 300, 300, 400, 400
212        let run_ends = Int32Array::from(vec![2i32, 5, 8, 10]);
213        let values = Int32Array::from(vec![100, 200, 300, 400]);
214        let arrow_run_array = RunArray::<Int32Type>::try_new(&run_ends, &values)?;
215
216        // Slice the array from index 1 to 7 (length 6)
217        // This should give us: 100, 200, 200, 200, 300, 300
218        let sliced_array = arrow_run_array.slice(1, 6);
219
220        // Convert the sliced array to Vortex
221        let vortex_array = decode_run_array(&sliced_array, false)?;
222        assert_arrays_eq!(
223            vortex_array,
224            buffer![100, 200, 200, 200, 300, 300].into_array(),
225            &mut SESSION.create_execution_ctx()
226        );
227        Ok(())
228    }
229
230    #[test]
231    fn test_sliced_arrow_run_array_with_nulls_to_vortex() -> VortexResult<()> {
232        // Create an Arrow RunArray with nullable values
233        // Run ends: [3, 6, 9, 12] means runs of length 3, 3, 3, 3
234        // Values: [Some(10), None, Some(30), Some(40)]
235        let run_ends = Int64Array::from(vec![3i64, 6, 9, 12]);
236        let values = Int64Array::from(vec![Some(10), None, Some(30), Some(40)]);
237        let arrow_run_array = RunArray::<Int64Type>::try_new(&run_ends, &values)?;
238
239        // Slice from index 4 to 10 (length 6)
240        // Original: 10, 10, 10, null, null, null, 30, 30, 30, 40, 40, 40
241        // Sliced:   null, null, 30, 30, 30, 40
242        let sliced_array = arrow_run_array.slice(4, 6);
243
244        // Convert to Vortex with nullable=true
245        let vortex_array = decode_run_array(&sliced_array, true)?;
246
247        assert_arrays_eq!(
248            vortex_array,
249            PrimitiveArray::from_option_iter([
250                None,
251                None,
252                Some(30i64),
253                Some(30),
254                Some(30),
255                Some(40),
256            ]),
257            &mut SESSION.create_execution_ctx()
258        );
259        Ok(())
260    }
261
262    #[test]
263    fn test_sliced_to_0_arrow_run_array_with_nulls_to_vortex() -> VortexResult<()> {
264        // Create an Arrow RunArray with nullable values
265        // Run ends: [3, 6, 9, 12] means runs of length 3, 3, 3, 3
266        // Values: [Some(10), None, Some(30), Some(40)]
267        let run_ends = Int64Array::from(vec![3i64, 6, 9, 12]);
268        let values = Int64Array::from(vec![Some(10), None, Some(30), Some(40)]);
269        let arrow_run_array = RunArray::<Int64Type>::try_new(&run_ends, &values)?;
270
271        // Slice from index 4 to 4 (length 0)
272        // Original: 10, 10, 10, null, null, null, 30, 30, 30, 40, 40, 40
273        // Sliced:   [ ]
274        let sliced_array = arrow_run_array.slice(4, 0);
275
276        // Convert to Vortex with nullable=true
277        let vortex_array = decode_run_array(&sliced_array, true)?;
278
279        // Verify properties
280        assert_eq!(vortex_array.len(), 0);
281        assert_eq!(
282            vortex_array.dtype(),
283            &DType::Primitive(PType::I64, Nullability::Nullable)
284        );
285        Ok(())
286    }
287
288    fn ree_type(ends: DataType, values_dtype: DataType) -> DataType {
289        DataType::RunEndEncoded(
290            Arc::new(Field::new("run_ends", ends, false)),
291            Arc::new(Field::new("values", values_dtype, true)),
292        )
293    }
294
295    fn execute(array: ArrayRef, dt: &DataType) -> VortexResult<arrow_array::ArrayRef> {
296        let field = Field::new("", dt.clone(), true);
297        SESSION
298            .arrow()
299            .execute_arrow(array, Some(&field), &mut SESSION.create_execution_ctx())
300    }
301
302    #[test]
303    fn test_roundtrip_arrow_to_vortex_to_arrow() -> VortexResult<()> {
304        let original = RunArray::<Int32Type>::try_new(
305            &Int32Array::from(vec![3i32, 5, 8]),
306            &Int32Array::from(vec![10, 20, 30]),
307        )?;
308        let vortex_array = decode_run_array(&original, false)?;
309        let target = ree_type(DataType::Int32, DataType::Int32);
310        let result = execute(vortex_array.into_array(), &target)?;
311
312        let expected = RunArray::<Int32Type>::try_new(
313            &Int32Array::from(vec![3, 5, 8]),
314            &Int32Array::from(vec![10, 20, 30]),
315        )?;
316        assert_eq!(result.as_ref(), &expected);
317        Ok(())
318    }
319
320    /// Slicing a RunEndArray and converting to Arrow REE must produce
321    /// correctly trimmed and adjusted run ends for both zero and non-zero offsets.
322    #[rstest]
323    #[case::nonzero_offset(
324        &[10i32, 10, 20, 20, 20, 30, 30],
325        1..5usize,
326        &[1i32, 4],
327        &[10i32, 20],
328    )]
329    #[case::zero_offset_excess_runs(
330        &[10i32, 10, 10, 20, 20, 30, 30, 30, 30, 30],
331        0..4usize,
332        &[3i32, 4],
333        &[10i32, 20],
334    )]
335    fn sliced_runend_to_arrow_ree(
336        #[case] input: &[i32],
337        #[case] slice_range: std::ops::Range<usize>,
338        #[case] expected_ends: &[i32],
339        #[case] expected_values: &[i32],
340    ) -> VortexResult<()> {
341        let mut ctx = SESSION.create_execution_ctx();
342        let array = RunEnd::encode(
343            PrimitiveArray::from_iter(input.iter().copied()).into_array(),
344            &mut ctx,
345        )?;
346        let sliced = array.into_array().slice(slice_range.clone())?;
347        let target = ree_type(DataType::Int32, DataType::Int32);
348        let result = execute(sliced, &target)?;
349
350        assert_eq!(result.len(), slice_range.len());
351        let expected = RunArray::<Int32Type>::try_new(
352            &Int32Array::from(expected_ends.to_vec()),
353            &Int32Array::from(expected_values.to_vec()),
354        )?;
355        assert_eq!(result.as_ref(), &expected);
356        Ok(())
357    }
358}
vortex_arrow/run_end_import.rs

vortex_arrow/
run_end_import.rs
