arrow_data/transform/
run.rs

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

arrow_data/transform/
run.rs
