arrow_data/transform/
run.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 super::{_MutableArrayData, ArrayData, Extend};
6use arrow_buffer::{ArrowNativeType, ToByteSlice};
7use arrow_schema::{ArrowError, DataType};
8use num_traits::CheckedAdd;
9
10/// Generic helper to get the last run end value from a run ends array
11fn get_last_run_end<T: ArrowNativeType>(run_ends_data: &super::MutableArrayData) -> T {
12    if run_ends_data.data.len == 0 {
13        T::default()
14    } else {
15        let typed_slice: &[T] = run_ends_data.data.buffer1.typed_data();
16        if typed_slice.len() >= run_ends_data.data.len {
17            typed_slice[run_ends_data.data.len - 1]
18        } else {
19            T::default()
20        }
21    }
22}
23
24/// Extends the `MutableArrayData` with null values.
25///
26/// For RunEndEncoded, this adds nulls by extending the run_ends array
27/// and values array appropriately.
28pub fn extend_nulls(mutable: &mut _MutableArrayData, len: usize) -> Result<(), ArrowError> {
29    if len == 0 {
30        return Ok(());
31    }
32
33    // For REE, we always need to add a value entry when adding a new run
34    // The values array should have one entry per run, not per logical element
35    mutable.child_data[1].try_extend_nulls(1)?;
36
37    // Determine the run end type from the data type
38    let run_end_type = if let DataType::RunEndEncoded(run_ends_field, _) = &mutable.data_type {
39        run_ends_field.data_type()
40    } else {
41        panic!("extend_nulls called on non-RunEndEncoded array");
42    };
43
44    // Use a macro to handle all run end types generically
45    macro_rules! extend_nulls_impl {
46        ($run_end_type:ty) => {{
47            let last_run_end = get_last_run_end::<$run_end_type>(&mutable.child_data[0]);
48            let len = <$run_end_type as ArrowNativeType>::from_usize(len);
49            let new_value = len
50                .and_then(|len| last_run_end.checked_add(len))
51                .ok_or_else(|| {
52                    ArrowError::InvalidArgumentError(
53                        "run end overflow when extending RunEndEncoded array: \
54                         use a larger run-end type (e.g. Int64 instead of Int32)"
55                            .to_string(),
56                    )
57                })?;
58            mutable.child_data[0]
59                .data
60                .buffer1
61                .try_extend_from_slice(new_value.to_byte_slice())
62                .map_err(|e| ArrowError::MemoryError(e.to_string()))?;
63        }};
64    }
65
66    // Apply the appropriate implementation based on run end type
67    match run_end_type {
68        DataType::Int16 => extend_nulls_impl!(i16),
69        DataType::Int32 => extend_nulls_impl!(i32),
70        DataType::Int64 => extend_nulls_impl!(i64),
71        _ => panic!("Invalid run end type for RunEndEncoded array: {run_end_type}"),
72    }
73
74    mutable.child_data[0].data.len += 1;
75    Ok(())
76}
77
78/// The run-ends bytes and optional values index range returned by [`build_extend_arrays`].
79type ExtendArrays = (Vec<u8>, Option<(usize, usize)>);
80
81/// Build run ends bytes and values range directly for batch processing
82fn build_extend_arrays<T: ArrowNativeType + std::ops::Add<Output = T> + CheckedAdd>(
83    typed_slice: &[T],
84    length: usize,
85    start: usize,
86    len: usize,
87    dest_last_run_end: T,
88) -> Result<ExtendArrays, ArrowError> {
89    let mut run_ends_bytes = Vec::new();
90    let mut values_range: Option<(usize, usize)> = None;
91    let end = start + len;
92    let mut prev_end = 0;
93    let mut current_run_end = dest_last_run_end;
94
95    for i in 0..length {
96        if i < typed_slice.len() {
97            let run_end = typed_slice[i].to_usize().unwrap();
98
99            if prev_end <= start && run_end > start {
100                let start_offset = start - prev_end;
101                let end_offset = if run_end >= end {
102                    end - prev_end
103                } else {
104                    run_end - prev_end
105                };
106                current_run_end = current_run_end
107                    .checked_add(&T::usize_as(end_offset - start_offset))
108                    .ok_or_else(|| {
109                        ArrowError::InvalidArgumentError(
110                            "run end overflow when extending RunEndEncoded array: \
111                         use a larger run-end type (e.g. Int64 instead of Int32)"
112                                .to_string(),
113                        )
114                    })?;
115                run_ends_bytes.extend_from_slice(current_run_end.to_byte_slice());
116
117                // Start the range
118                values_range = Some((i, i + 1));
119            } else if prev_end >= start && run_end <= end {
120                current_run_end = current_run_end
121                    .checked_add(&T::usize_as(run_end - prev_end))
122                    .ok_or_else(|| {
123                        ArrowError::InvalidArgumentError(
124                            "run end overflow when extending RunEndEncoded array: \
125                         use a larger run-end type (e.g. Int64 instead of Int32)"
126                                .to_string(),
127                        )
128                    })?;
129                run_ends_bytes.extend_from_slice(current_run_end.to_byte_slice());
130
131                // Extend the range
132                values_range = Some((values_range.expect("Unreachable: values_range cannot be None when prev_end >= start && run_end <= end. \
133                           If prev_end >= start and run_end > prev_end (required for valid runs), then run_end > start, \
134                           which means the first condition (prev_end <= start && run_end > start) would have been true \
135                           and already set values_range to Some.").0, i + 1));
136            } else if prev_end < end && run_end >= end {
137                current_run_end = current_run_end
138                    .checked_add(&T::usize_as(end - prev_end))
139                    .ok_or_else(|| {
140                        ArrowError::InvalidArgumentError(
141                            "run end overflow when extending RunEndEncoded array: \
142                         use a larger run-end type (e.g. Int64 instead of Int32)"
143                                .to_string(),
144                        )
145                    })?;
146                run_ends_bytes.extend_from_slice(current_run_end.to_byte_slice());
147
148                // Extend the range and break
149                values_range = Some((values_range.expect("Unreachable: values_range cannot be None when prev_end < end && run_end >= end. \
150                           Due to sequential processing and monotonic prev_end advancement, if we reach a run \
151                           that spans beyond the slice end (run_end >= end), at least one previous condition \
152                           must have matched first to set values_range. Either the first condition matched when \
153                           the slice started (prev_end <= start && run_end > start), or the second condition \
154                           matched for runs within the slice (prev_end >= start && run_end <= end).").0, i + 1));
155                break;
156            }
157
158            prev_end = run_end;
159            if prev_end >= end {
160                break;
161            }
162        } else {
163            break;
164        }
165    }
166    Ok((run_ends_bytes, values_range))
167}
168
169/// Process extends using batch operations
170fn process_extends_batch<T: ArrowNativeType>(
171    mutable: &mut _MutableArrayData,
172    source_array_idx: usize,
173    run_ends_bytes: Vec<u8>,
174    values_range: Option<(usize, usize)>,
175) -> Result<(), ArrowError> {
176    if run_ends_bytes.is_empty() {
177        return Ok(());
178    }
179
180    // Batch extend the run_ends array with all bytes at once
181    mutable.child_data[0]
182        .data
183        .buffer1
184        .extend_from_slice(&run_ends_bytes);
185    mutable.child_data[0].data.len += run_ends_bytes.len() / std::mem::size_of::<T>();
186
187    // Batch extend the values array using the range
188    let (start_idx, end_idx) =
189        values_range.expect("values_range should be Some if run_ends_bytes is not empty");
190    mutable.child_data[1].try_extend(source_array_idx, start_idx, end_idx)
191}
192
193/// Returns a function that extends the run encoded array.
194///
195/// It finds the physical indices in the source array that correspond to the logical range to copy, and adjusts the runs to the logical indices of the array to extend. The values are copied from the source array to the destination array verbatim.
196pub fn build_extend(array: &ArrayData) -> Extend<'_> {
197    Box::new(
198        move |mutable: &mut _MutableArrayData, array_idx: usize, start: usize, len: usize| {
199            if len == 0 {
200                return Ok(());
201            }
202
203            // We need to analyze the source array's run structure
204            let source_run_ends = &array.child_data()[0];
205            let source_buffer_offset = source_run_ends.offset();
206            let source_buffer_len = source_run_ends.len();
207            let source_buffer = &source_run_ends.buffers()[0];
208
209            // Get the run end type from the mutable array
210            let dest_run_end_type =
211                if let DataType::RunEndEncoded(run_ends_field, _) = &mutable.data_type {
212                    run_ends_field.data_type()
213                } else {
214                    panic!("extend called on non-RunEndEncoded mutable array");
215                };
216
217            // Build run ends and values indices directly for batch processing
218            macro_rules! build_and_process_impl {
219                ($run_end_type:ty) => {{
220                    let source_buffer = &source_buffer.typed_data()
221                        [source_buffer_offset..source_buffer_offset + source_buffer_len];
222                    let dest_last_run_end =
223                        get_last_run_end::<$run_end_type>(&mutable.child_data[0]);
224                    let (run_ends_bytes, values_range) = build_extend_arrays::<$run_end_type>(
225                        source_buffer,
226                        source_run_ends.len(),
227                        start + array.offset(),
228                        len,
229                        dest_last_run_end,
230                    )?;
231                    process_extends_batch::<$run_end_type>(
232                        mutable,
233                        array_idx,
234                        run_ends_bytes,
235                        values_range,
236                    )?;
237                }};
238            }
239
240            match dest_run_end_type {
241                DataType::Int16 => build_and_process_impl!(i16),
242                DataType::Int32 => build_and_process_impl!(i32),
243                DataType::Int64 => build_and_process_impl!(i64),
244                _ => panic!("Invalid run end type for RunEndEncoded array: {dest_run_end_type}"),
245            }
246            Ok(())
247        },
248    )
249}
250
251#[cfg(test)]
252mod tests {
253    use super::*;
254    use crate::transform::MutableArrayData;
255    use crate::{ArrayData, ArrayDataBuilder};
256    use arrow_buffer::Buffer;
257    use arrow_schema::{DataType, Field};
258    use std::sync::Arc;
259
260    fn create_run_array_data(run_ends: Vec<i32>, values: ArrayData) -> ArrayData {
261        let run_ends_field = Arc::new(Field::new(
262            Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
263            DataType::Int32,
264            false,
265        ));
266        let values_field = Arc::new(Field::new(
267            Field::REE_VALUES_FIELD_DEFAULT_NAME,
268            values.data_type().clone(),
269            true,
270        ));
271        let data_type = DataType::RunEndEncoded(run_ends_field, values_field);
272
273        let last_run_end = if run_ends.is_empty() {
274            0
275        } else {
276            run_ends[run_ends.len() - 1] as usize
277        };
278
279        let run_ends_buffer = Buffer::from_vec(run_ends);
280        let run_ends_data = ArrayDataBuilder::new(DataType::Int32)
281            .len(run_ends_buffer.len() / std::mem::size_of::<i32>())
282            .add_buffer(run_ends_buffer)
283            .build()
284            .unwrap();
285
286        ArrayDataBuilder::new(data_type)
287            .len(last_run_end)
288            .add_child_data(run_ends_data)
289            .add_child_data(values)
290            .build()
291            .unwrap()
292    }
293
294    fn create_run_array_data_int16(run_ends: Vec<i16>, values: ArrayData) -> ArrayData {
295        let run_ends_field = Arc::new(Field::new(
296            Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
297            DataType::Int16,
298            false,
299        ));
300        let values_field = Arc::new(Field::new(
301            Field::REE_VALUES_FIELD_DEFAULT_NAME,
302            values.data_type().clone(),
303            true,
304        ));
305        let data_type = DataType::RunEndEncoded(run_ends_field, values_field);
306
307        let last_run_end = if run_ends.is_empty() {
308            0
309        } else {
310            run_ends[run_ends.len() - 1] as usize
311        };
312
313        let run_ends_buffer = Buffer::from_vec(run_ends);
314        let run_ends_data = ArrayDataBuilder::new(DataType::Int16)
315            .len(run_ends_buffer.len() / std::mem::size_of::<i16>())
316            .add_buffer(run_ends_buffer)
317            .build()
318            .unwrap();
319
320        ArrayDataBuilder::new(data_type)
321            .len(last_run_end)
322            .add_child_data(run_ends_data)
323            .add_child_data(values)
324            .build()
325            .unwrap()
326    }
327
328    fn create_run_array_data_int64(run_ends: Vec<i64>, values: ArrayData) -> ArrayData {
329        let run_ends_field = Arc::new(Field::new(
330            Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
331            DataType::Int64,
332            false,
333        ));
334        let values_field = Arc::new(Field::new(
335            Field::REE_VALUES_FIELD_DEFAULT_NAME,
336            values.data_type().clone(),
337            true,
338        ));
339        let data_type = DataType::RunEndEncoded(run_ends_field, values_field);
340
341        let last_run_end = if run_ends.is_empty() {
342            0
343        } else {
344            run_ends[run_ends.len() - 1] as usize
345        };
346
347        let run_ends_buffer = Buffer::from_vec(run_ends);
348        let run_ends_data = ArrayDataBuilder::new(DataType::Int64)
349            .len(run_ends_buffer.len() / std::mem::size_of::<i64>())
350            .add_buffer(run_ends_buffer)
351            .build()
352            .unwrap();
353
354        ArrayDataBuilder::new(data_type)
355            .len(last_run_end)
356            .add_child_data(run_ends_data)
357            .add_child_data(values)
358            .build()
359            .unwrap()
360    }
361
362    fn create_int32_array_data(values: Vec<i32>) -> ArrayData {
363        let buffer = Buffer::from_vec(values);
364        ArrayDataBuilder::new(DataType::Int32)
365            .len(buffer.len() / std::mem::size_of::<i32>())
366            .add_buffer(buffer)
367            .build()
368            .unwrap()
369    }
370
371    fn create_string_dict_array_data(values: Vec<&str>, dict_values: Vec<&str>) -> ArrayData {
372        // Create dictionary values (strings)
373        let dict_offsets: Vec<i32> = dict_values
374            .iter()
375            .scan(0i32, |acc, s| {
376                let offset = *acc;
377                *acc += s.len() as i32;
378                Some(offset)
379            })
380            .chain(std::iter::once(
381                dict_values.iter().map(|s| s.len()).sum::<usize>() as i32,
382            ))
383            .collect();
384
385        let dict_data: Vec<u8> = dict_values.iter().flat_map(|s| s.bytes()).collect();
386
387        let dict_array = ArrayDataBuilder::new(DataType::Utf8)
388            .len(dict_values.len())
389            .add_buffer(Buffer::from_vec(dict_offsets))
390            .add_buffer(Buffer::from_vec(dict_data))
391            .build()
392            .unwrap();
393
394        // Create keys array
395        let keys: Vec<i32> = values
396            .iter()
397            .map(|v| dict_values.iter().position(|d| d == v).unwrap() as i32)
398            .collect();
399
400        // Create dictionary array
401        let dict_type = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
402
403        ArrayDataBuilder::new(dict_type)
404            .len(values.len())
405            .add_buffer(Buffer::from_vec(keys))
406            .add_child_data(dict_array)
407            .build()
408            .unwrap()
409    }
410
411    #[test]
412    fn test_extend_nulls_int32() {
413        // Create values array with one value
414        let values = create_int32_array_data(vec![42]);
415
416        // Create REE array with Int32 run ends
417        let ree_array = create_run_array_data(vec![5], values);
418
419        let mut mutable = MutableArrayData::new(vec![&ree_array], true, 10);
420
421        mutable.try_extend_nulls(3).unwrap();
422        mutable.try_extend(0, 0, 5).unwrap();
423        mutable.try_extend_nulls(3).unwrap();
424
425        // Verify the run ends were extended correctly
426        let result = mutable.freeze();
427        let run_ends_buffer = &result.child_data()[0].buffers()[0];
428        let run_ends_slice = run_ends_buffer.as_slice();
429
430        // Should have three run ends now
431        assert_eq!(result.child_data()[0].len(), 3);
432        let first_run_end = i32::from_ne_bytes(run_ends_slice[0..4].try_into().unwrap());
433        let second_run_end = i32::from_ne_bytes(run_ends_slice[4..8].try_into().unwrap());
434        let third_run_end = i32::from_ne_bytes(run_ends_slice[8..12].try_into().unwrap());
435        assert_eq!(first_run_end, 3);
436        assert_eq!(second_run_end, 8);
437        assert_eq!(third_run_end, 11);
438
439        // Verify the values array was extended correctly
440        assert_eq!(result.child_data()[1].len(), 3); // Should match run ends length
441        let values_buffer = &result.child_data()[1].buffers()[0];
442        let values_slice = values_buffer.as_slice();
443
444        // Check the values in the buffer
445        let second_value = i32::from_ne_bytes(values_slice[4..8].try_into().unwrap());
446
447        // Second value should be the original value from the source array
448        assert_eq!(second_value, 42);
449
450        // Verify the validity buffer shows the correct null pattern
451        let values_array = &result.child_data()[1];
452        // First value should be null
453        assert!(values_array.is_null(0));
454        // Second value should be valid
455        assert!(values_array.is_valid(1));
456        // Third value should be null
457        assert!(values_array.is_null(2));
458    }
459
460    #[test]
461    fn test_extend_nulls_int16() {
462        // Create values array with one value
463        let values = create_int32_array_data(vec![42]);
464
465        // Create REE array with Int16 run ends
466        let ree_array = create_run_array_data_int16(vec![5i16], values);
467
468        let mut mutable = MutableArrayData::new(vec![&ree_array], true, 10);
469
470        // First, we need to copy the existing data
471        mutable.try_extend(0, 0, 5).unwrap();
472
473        // Then add nulls
474        mutable.try_extend_nulls(3).unwrap();
475
476        // Verify the run ends were extended correctly
477        let result = mutable.freeze();
478        let run_ends_buffer = &result.child_data()[0].buffers()[0];
479        let run_ends_slice = run_ends_buffer.as_slice();
480
481        // Should have two run ends now: original 5 and new 8 (5 + 3)
482        assert_eq!(result.child_data()[0].len(), 2);
483        let first_run_end = i16::from_ne_bytes(run_ends_slice[0..2].try_into().unwrap());
484        let second_run_end = i16::from_ne_bytes(run_ends_slice[2..4].try_into().unwrap());
485        assert_eq!(first_run_end, 5);
486        assert_eq!(second_run_end, 8);
487    }
488
489    #[test]
490    fn test_extend_nulls_int64() {
491        // Create values array with one value
492        let values = create_int32_array_data(vec![42]);
493
494        // Create REE array with Int64 run ends
495        let ree_array = create_run_array_data_int64(vec![5i64], values);
496
497        let mut mutable = MutableArrayData::new(vec![&ree_array], true, 10);
498
499        // First, we need to copy the existing data
500        mutable.try_extend(0, 0, 5).unwrap();
501
502        // Then add nulls
503        mutable.try_extend_nulls(3).unwrap();
504
505        // Verify the run ends were extended correctly
506        let result = mutable.freeze();
507        let run_ends_buffer = &result.child_data()[0].buffers()[0];
508        let run_ends_slice = run_ends_buffer.as_slice();
509
510        // Should have two run ends now: original 5 and new 8 (5 + 3)
511        assert_eq!(result.child_data()[0].len(), 2);
512        let first_run_end = i64::from_ne_bytes(run_ends_slice[0..8].try_into().unwrap());
513        let second_run_end = i64::from_ne_bytes(run_ends_slice[8..16].try_into().unwrap());
514        assert_eq!(first_run_end, 5);
515        assert_eq!(second_run_end, 8);
516    }
517
518    #[test]
519    fn test_extend_int32() {
520        // Create a simple REE array with Int32 run ends
521        let values = create_int32_array_data(vec![10, 20]);
522
523        // Array: [10, 10, 20, 20, 20] (run_ends = [2, 5])
524        let ree_array = create_run_array_data(vec![2, 5], values);
525
526        let mut mutable = MutableArrayData::new(vec![&ree_array], false, 10);
527
528        // Extend the entire array
529        mutable.try_extend(0, 0, 5).unwrap();
530
531        let result = mutable.freeze();
532
533        // Should have extended correctly
534        assert_eq!(result.len(), 5); // All 5 elements
535
536        // Basic validation that we have the right structure
537        assert!(!result.child_data()[0].is_empty()); // Should have at least one run
538        assert_eq!(result.child_data()[0].len(), result.child_data()[1].len()); // run_ends and values should have same length
539    }
540
541    #[test]
542    fn test_extend_empty() {
543        let values = create_int32_array_data(vec![]);
544        let ree_array = create_run_array_data(vec![], values);
545
546        let mut mutable = MutableArrayData::new(vec![&ree_array], false, 10);
547        mutable.try_extend(0, 0, 0).unwrap();
548
549        let result = mutable.freeze();
550        assert_eq!(result.len(), 0);
551        assert_eq!(result.child_data()[0].len(), 0);
552    }
553
554    #[test]
555    fn test_build_extend_arrays_int16() {
556        let buffer = Buffer::from_vec(vec![3i16, 5i16, 8i16]);
557        let (run_ends_bytes, values_range) =
558            build_extend_arrays::<i16>(buffer.typed_data(), 3, 2, 4, 0i16).unwrap();
559
560        // Logical array: [A, A, A, B, B, C, C, C]
561        // Requesting indices 2-6 should give us:
562        // - Part of first run (index 2) -> length 1
563        // - All of second run -> length 2
564        // - Part of third run -> length 1
565        // Total length = 4, so run ends should be [1, 3, 4]
566        assert_eq!(run_ends_bytes.len(), 3 * std::mem::size_of::<i16>());
567        assert_eq!(values_range, Some((0, 3)));
568
569        // Verify the bytes represent [1i16, 3i16, 4i16]
570        let expected_bytes = [1i16, 3i16, 4i16]
571            .iter()
572            .flat_map(|&val| val.to_ne_bytes())
573            .collect::<Vec<u8>>();
574        assert_eq!(run_ends_bytes, expected_bytes);
575    }
576
577    #[test]
578    fn test_build_extend_arrays_int64() {
579        let buffer = Buffer::from_vec(vec![3i64, 5i64, 8i64]);
580        let (run_ends_bytes, values_range) =
581            build_extend_arrays::<i64>(buffer.typed_data(), 3, 2, 4, 0i64).unwrap();
582
583        // Same logic as above but with i64
584        assert_eq!(run_ends_bytes.len(), 3 * std::mem::size_of::<i64>());
585        assert_eq!(values_range, Some((0, 3)));
586
587        // Verify the bytes represent [1i64, 3i64, 4i64]
588        let expected_bytes = [1i64, 3i64, 4i64]
589            .iter()
590            .flat_map(|&val| val.to_ne_bytes())
591            .collect::<Vec<u8>>();
592        assert_eq!(run_ends_bytes, expected_bytes);
593    }
594
595    #[test]
596    fn test_extend_string_dict() {
597        // Create a dictionary array with string values: ["hello", "world"]
598        let dict_values = vec!["hello", "world"];
599        let values = create_string_dict_array_data(vec!["hello", "world"], dict_values);
600
601        // Create REE array: [hello, hello, world, world, world] (run_ends = [2, 5])
602        let ree_array = create_run_array_data(vec![2, 5], values);
603
604        let mut mutable = MutableArrayData::new(vec![&ree_array], false, 10);
605
606        // Extend the entire array
607        mutable.try_extend(0, 0, 5).unwrap();
608
609        let result = mutable.freeze();
610
611        // Should have extended correctly
612        assert_eq!(result.len(), 5); // All 5 elements
613
614        // Basic validation that we have the right structure
615        assert!(!result.child_data()[0].is_empty()); // Should have at least one run
616        assert_eq!(result.child_data()[0].len(), result.child_data()[1].len()); // run_ends and values should have same length
617
618        // Should have 2 runs since we have 2 different values
619        assert_eq!(result.child_data()[0].len(), 2);
620        assert_eq!(result.child_data()[1].len(), 2);
621    }
622
623    #[test]
624    fn test_extend_nulls_overflow_i16() {
625        let values = create_int32_array_data(vec![42]);
626        // Start with run end close to max to set up overflow condition
627        let ree_array = create_run_array_data_int16(vec![5], values);
628        let mut mutable = MutableArrayData::new(vec![&ree_array], true, 10);
629
630        // Extend the original data first to initialize state
631        mutable.try_extend(0, 0, 5_usize).unwrap();
632
633        // This should return an error: i16::MAX + 5 > i16::MAX
634        let err = mutable.try_extend_nulls(i16::MAX as usize).unwrap_err();
635        assert!(
636            err.to_string().contains("run end overflow"),
637            "unexpected error: {err}"
638        );
639    }
640
641    #[test]
642    fn test_extend_nulls_overflow_i32() {
643        let values = create_int32_array_data(vec![42]);
644        // Start with run end close to max to set up overflow condition
645        let ree_array = create_run_array_data(vec![10], values);
646        let mut mutable = MutableArrayData::new(vec![&ree_array], true, 10);
647
648        // Extend the original data first to initialize state
649        mutable.try_extend(0, 0, 10_usize).unwrap();
650
651        // This should return an error: (i32::MAX - 10) + 20 > i32::MAX
652        let err = mutable.try_extend_nulls(i32::MAX as usize).unwrap_err();
653        assert!(
654            err.to_string().contains("run end overflow"),
655            "unexpected error: {err}"
656        );
657    }
658
659    #[test]
660    fn test_build_extend_overflow_i16() {
661        // Create a source array with small run that will cause overflow when added
662        let values = create_int32_array_data(vec![10]);
663        let source_array = create_run_array_data_int16(vec![20], values);
664
665        // Create a destination array with run end close to max
666        let dest_values = create_int32_array_data(vec![42]);
667        let dest_array = create_run_array_data_int16(vec![i16::MAX - 5], dest_values);
668
669        let mut mutable = MutableArrayData::new(vec![&source_array, &dest_array], false, 10);
670
671        // First extend the destination array to set up state
672        mutable.try_extend(1, 0, (i16::MAX - 5) as usize).unwrap();
673
674        // This should return an error: (i16::MAX - 5) + 20 > i16::MAX
675        let err = mutable.try_extend(0, 0, 20).unwrap_err();
676        assert!(
677            err.to_string().contains("run end overflow"),
678            "unexpected error: {err}"
679        );
680    }
681
682    #[test]
683    fn test_build_extend_overflow_i32() {
684        // Create a source array with small run that will cause overflow when added
685        let values = create_int32_array_data(vec![10]);
686        let source_array = create_run_array_data(vec![100], values);
687
688        // Create a destination array with run end close to max
689        let dest_values = create_int32_array_data(vec![42]);
690        let dest_array = create_run_array_data(vec![i32::MAX - 50], dest_values);
691
692        let mut mutable = MutableArrayData::new(vec![&source_array, &dest_array], false, 10);
693
694        // First extend the destination array to set up state
695        mutable.try_extend(1, 0, (i32::MAX - 50) as usize).unwrap();
696
697        // This should return an error: (i32::MAX - 50) + 100 > i32::MAX
698        let err = mutable.try_extend(0, 0, 100).unwrap_err();
699        assert!(
700            err.to_string().contains("run end overflow"),
701            "unexpected error: {err}"
702        );
703    }
704
705    #[test]
706    fn test_extend_nulls_overflow() {
707        let values = create_int32_array_data(vec![10]);
708        let array = create_run_array_data_int16(vec![100], values);
709
710        let mut mutable = MutableArrayData::new(vec![&array], true, 10);
711        let err = mutable.try_extend_nulls(i16::MAX as usize + 1).unwrap_err();
712        assert_eq!(
713            err.to_string(),
714            "Invalid argument error: run end overflow when extending RunEndEncoded array: use a larger run-end type (e.g. Int64 instead of Int32)"
715        );
716    }
717
718    #[test]
719    fn test_extend_accounts_for_offset() {
720        let dt = DataType::RunEndEncoded(
721            Arc::new(Field::new(
722                Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
723                DataType::Int32,
724                false,
725            )),
726            Arc::new(Field::new(
727                Field::REE_VALUES_FIELD_DEFAULT_NAME,
728                DataType::Float32,
729                true,
730            )),
731        );
732
733        let runs = ArrayData::builder(DataType::Int32)
734            .len(2)
735            .offset(2)
736            .add_buffer(vec![-10_i32, -20, 1_i32, 4, 0, 100].into())
737            .build()
738            .unwrap();
739        let values = ArrayData::builder(DataType::Float32)
740            .len(2)
741            .add_buffer(vec![1.5, 2.5].into())
742            .build()
743            .unwrap();
744
745        let source_data = ArrayDataBuilder::new(dt)
746            .len(3)
747            .child_data(vec![runs, values])
748            .build()
749            .unwrap();
750
751        let mut mutable = MutableArrayData::new(vec![&source_data], true, 10);
752        mutable.try_extend(0, 0, 1).unwrap();
753        let data = mutable.freeze();
754        assert_eq!(data, source_data.slice(0, 1));
755    }
756}
arrow_data/transform/run.rs

arrow_data/transform/
run.rs
