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("run_ends", DataType::Int32, false));
273        let values_field = Arc::new(Field::new("values", values.data_type().clone(), true));
274        let data_type = DataType::RunEndEncoded(run_ends_field, values_field);
275
276        let last_run_end = if run_ends.is_empty() {
277            0
278        } else {
279            run_ends[run_ends.len() - 1] as usize
280        };
281
282        let run_ends_buffer = Buffer::from_vec(run_ends);
283        let run_ends_data = ArrayDataBuilder::new(DataType::Int32)
284            .len(run_ends_buffer.len() / std::mem::size_of::<i32>())
285            .add_buffer(run_ends_buffer)
286            .build()
287            .unwrap();
288
289        ArrayDataBuilder::new(data_type)
290            .len(last_run_end)
291            .add_child_data(run_ends_data)
292            .add_child_data(values)
293            .build()
294            .unwrap()
295    }
296
297    fn create_run_array_data_int16(run_ends: Vec<i16>, values: ArrayData) -> ArrayData {
298        let run_ends_field = Arc::new(Field::new("run_ends", DataType::Int16, false));
299        let values_field = Arc::new(Field::new("values", values.data_type().clone(), true));
300        let data_type = DataType::RunEndEncoded(run_ends_field, values_field);
301
302        let last_run_end = if run_ends.is_empty() {
303            0
304        } else {
305            run_ends[run_ends.len() - 1] as usize
306        };
307
308        let run_ends_buffer = Buffer::from_vec(run_ends);
309        let run_ends_data = ArrayDataBuilder::new(DataType::Int16)
310            .len(run_ends_buffer.len() / std::mem::size_of::<i16>())
311            .add_buffer(run_ends_buffer)
312            .build()
313            .unwrap();
314
315        ArrayDataBuilder::new(data_type)
316            .len(last_run_end)
317            .add_child_data(run_ends_data)
318            .add_child_data(values)
319            .build()
320            .unwrap()
321    }
322
323    fn create_run_array_data_int64(run_ends: Vec<i64>, values: ArrayData) -> ArrayData {
324        let run_ends_field = Arc::new(Field::new("run_ends", DataType::Int64, false));
325        let values_field = Arc::new(Field::new("values", values.data_type().clone(), true));
326        let data_type = DataType::RunEndEncoded(run_ends_field, values_field);
327
328        let last_run_end = if run_ends.is_empty() {
329            0
330        } else {
331            run_ends[run_ends.len() - 1] as usize
332        };
333
334        let run_ends_buffer = Buffer::from_vec(run_ends);
335        let run_ends_data = ArrayDataBuilder::new(DataType::Int64)
336            .len(run_ends_buffer.len() / std::mem::size_of::<i64>())
337            .add_buffer(run_ends_buffer)
338            .build()
339            .unwrap();
340
341        ArrayDataBuilder::new(data_type)
342            .len(last_run_end)
343            .add_child_data(run_ends_data)
344            .add_child_data(values)
345            .build()
346            .unwrap()
347    }
348
349    fn create_int32_array_data(values: Vec<i32>) -> ArrayData {
350        let buffer = Buffer::from_vec(values);
351        ArrayDataBuilder::new(DataType::Int32)
352            .len(buffer.len() / std::mem::size_of::<i32>())
353            .add_buffer(buffer)
354            .build()
355            .unwrap()
356    }
357
358    fn create_string_dict_array_data(values: Vec<&str>, dict_values: Vec<&str>) -> ArrayData {
359        // Create dictionary values (strings)
360        let dict_offsets: Vec<i32> = dict_values
361            .iter()
362            .scan(0i32, |acc, s| {
363                let offset = *acc;
364                *acc += s.len() as i32;
365                Some(offset)
366            })
367            .chain(std::iter::once(
368                dict_values.iter().map(|s| s.len()).sum::<usize>() as i32,
369            ))
370            .collect();
371
372        let dict_data: Vec<u8> = dict_values.iter().flat_map(|s| s.bytes()).collect();
373
374        let dict_array = ArrayDataBuilder::new(DataType::Utf8)
375            .len(dict_values.len())
376            .add_buffer(Buffer::from_vec(dict_offsets))
377            .add_buffer(Buffer::from_vec(dict_data))
378            .build()
379            .unwrap();
380
381        // Create keys array
382        let keys: Vec<i32> = values
383            .iter()
384            .map(|v| dict_values.iter().position(|d| d == v).unwrap() as i32)
385            .collect();
386
387        // Create dictionary array
388        let dict_type = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
389
390        ArrayDataBuilder::new(dict_type)
391            .len(values.len())
392            .add_buffer(Buffer::from_vec(keys))
393            .add_child_data(dict_array)
394            .build()
395            .unwrap()
396    }
397
398    #[test]
399    fn test_extend_nulls_int32() {
400        // Create values array with one value
401        let values = create_int32_array_data(vec![42]);
402
403        // Create REE array with Int32 run ends
404        let ree_array = create_run_array_data(vec![5], values);
405
406        let mut mutable = MutableArrayData::new(vec![&ree_array], true, 10);
407
408        mutable.try_extend_nulls(3).unwrap();
409        mutable.try_extend(0, 0, 5).unwrap();
410        mutable.try_extend_nulls(3).unwrap();
411
412        // Verify the run ends were extended correctly
413        let result = mutable.freeze();
414        let run_ends_buffer = &result.child_data()[0].buffers()[0];
415        let run_ends_slice = run_ends_buffer.as_slice();
416
417        // Should have three run ends now
418        assert_eq!(result.child_data()[0].len(), 3);
419        let first_run_end = i32::from_ne_bytes(run_ends_slice[0..4].try_into().unwrap());
420        let second_run_end = i32::from_ne_bytes(run_ends_slice[4..8].try_into().unwrap());
421        let third_run_end = i32::from_ne_bytes(run_ends_slice[8..12].try_into().unwrap());
422        assert_eq!(first_run_end, 3);
423        assert_eq!(second_run_end, 8);
424        assert_eq!(third_run_end, 11);
425
426        // Verify the values array was extended correctly
427        assert_eq!(result.child_data()[1].len(), 3); // Should match run ends length
428        let values_buffer = &result.child_data()[1].buffers()[0];
429        let values_slice = values_buffer.as_slice();
430
431        // Check the values in the buffer
432        let second_value = i32::from_ne_bytes(values_slice[4..8].try_into().unwrap());
433
434        // Second value should be the original value from the source array
435        assert_eq!(second_value, 42);
436
437        // Verify the validity buffer shows the correct null pattern
438        let values_array = &result.child_data()[1];
439        // First value should be null
440        assert!(values_array.is_null(0));
441        // Second value should be valid
442        assert!(values_array.is_valid(1));
443        // Third value should be null
444        assert!(values_array.is_null(2));
445    }
446
447    #[test]
448    fn test_extend_nulls_int16() {
449        // Create values array with one value
450        let values = create_int32_array_data(vec![42]);
451
452        // Create REE array with Int16 run ends
453        let ree_array = create_run_array_data_int16(vec![5i16], values);
454
455        let mut mutable = MutableArrayData::new(vec![&ree_array], true, 10);
456
457        // First, we need to copy the existing data
458        mutable.try_extend(0, 0, 5).unwrap();
459
460        // Then add nulls
461        mutable.try_extend_nulls(3).unwrap();
462
463        // Verify the run ends were extended correctly
464        let result = mutable.freeze();
465        let run_ends_buffer = &result.child_data()[0].buffers()[0];
466        let run_ends_slice = run_ends_buffer.as_slice();
467
468        // Should have two run ends now: original 5 and new 8 (5 + 3)
469        assert_eq!(result.child_data()[0].len(), 2);
470        let first_run_end = i16::from_ne_bytes(run_ends_slice[0..2].try_into().unwrap());
471        let second_run_end = i16::from_ne_bytes(run_ends_slice[2..4].try_into().unwrap());
472        assert_eq!(first_run_end, 5);
473        assert_eq!(second_run_end, 8);
474    }
475
476    #[test]
477    fn test_extend_nulls_int64() {
478        // Create values array with one value
479        let values = create_int32_array_data(vec![42]);
480
481        // Create REE array with Int64 run ends
482        let ree_array = create_run_array_data_int64(vec![5i64], values);
483
484        let mut mutable = MutableArrayData::new(vec![&ree_array], true, 10);
485
486        // First, we need to copy the existing data
487        mutable.try_extend(0, 0, 5).unwrap();
488
489        // Then add nulls
490        mutable.try_extend_nulls(3).unwrap();
491
492        // Verify the run ends were extended correctly
493        let result = mutable.freeze();
494        let run_ends_buffer = &result.child_data()[0].buffers()[0];
495        let run_ends_slice = run_ends_buffer.as_slice();
496
497        // Should have two run ends now: original 5 and new 8 (5 + 3)
498        assert_eq!(result.child_data()[0].len(), 2);
499        let first_run_end = i64::from_ne_bytes(run_ends_slice[0..8].try_into().unwrap());
500        let second_run_end = i64::from_ne_bytes(run_ends_slice[8..16].try_into().unwrap());
501        assert_eq!(first_run_end, 5);
502        assert_eq!(second_run_end, 8);
503    }
504
505    #[test]
506    fn test_extend_int32() {
507        // Create a simple REE array with Int32 run ends
508        let values = create_int32_array_data(vec![10, 20]);
509
510        // Array: [10, 10, 20, 20, 20] (run_ends = [2, 5])
511        let ree_array = create_run_array_data(vec![2, 5], values);
512
513        let mut mutable = MutableArrayData::new(vec![&ree_array], false, 10);
514
515        // Extend the entire array
516        mutable.try_extend(0, 0, 5).unwrap();
517
518        let result = mutable.freeze();
519
520        // Should have extended correctly
521        assert_eq!(result.len(), 5); // All 5 elements
522
523        // Basic validation that we have the right structure
524        assert!(!result.child_data()[0].is_empty()); // Should have at least one run
525        assert_eq!(result.child_data()[0].len(), result.child_data()[1].len()); // run_ends and values should have same length
526    }
527
528    #[test]
529    fn test_extend_empty() {
530        let values = create_int32_array_data(vec![]);
531        let ree_array = create_run_array_data(vec![], values);
532
533        let mut mutable = MutableArrayData::new(vec![&ree_array], false, 10);
534        mutable.try_extend(0, 0, 0).unwrap();
535
536        let result = mutable.freeze();
537        assert_eq!(result.len(), 0);
538        assert_eq!(result.child_data()[0].len(), 0);
539    }
540
541    #[test]
542    fn test_build_extend_arrays_int16() {
543        let buffer = Buffer::from_vec(vec![3i16, 5i16, 8i16]);
544        let (run_ends_bytes, values_range) =
545            build_extend_arrays::<i16>(&buffer, 3, 2, 4, 0i16).unwrap();
546
547        // Logical array: [A, A, A, B, B, C, C, C]
548        // Requesting indices 2-6 should give us:
549        // - Part of first run (index 2) -> length 1
550        // - All of second run -> length 2
551        // - Part of third run -> length 1
552        // Total length = 4, so run ends should be [1, 3, 4]
553        assert_eq!(run_ends_bytes.len(), 3 * std::mem::size_of::<i16>());
554        assert_eq!(values_range, Some((0, 3)));
555
556        // Verify the bytes represent [1i16, 3i16, 4i16]
557        let expected_bytes = [1i16, 3i16, 4i16]
558            .iter()
559            .flat_map(|&val| val.to_ne_bytes())
560            .collect::<Vec<u8>>();
561        assert_eq!(run_ends_bytes, expected_bytes);
562    }
563
564    #[test]
565    fn test_build_extend_arrays_int64() {
566        let buffer = Buffer::from_vec(vec![3i64, 5i64, 8i64]);
567        let (run_ends_bytes, values_range) =
568            build_extend_arrays::<i64>(&buffer, 3, 2, 4, 0i64).unwrap();
569
570        // Same logic as above but with i64
571        assert_eq!(run_ends_bytes.len(), 3 * std::mem::size_of::<i64>());
572        assert_eq!(values_range, Some((0, 3)));
573
574        // Verify the bytes represent [1i64, 3i64, 4i64]
575        let expected_bytes = [1i64, 3i64, 4i64]
576            .iter()
577            .flat_map(|&val| val.to_ne_bytes())
578            .collect::<Vec<u8>>();
579        assert_eq!(run_ends_bytes, expected_bytes);
580    }
581
582    #[test]
583    fn test_extend_string_dict() {
584        // Create a dictionary array with string values: ["hello", "world"]
585        let dict_values = vec!["hello", "world"];
586        let values = create_string_dict_array_data(vec!["hello", "world"], dict_values);
587
588        // Create REE array: [hello, hello, world, world, world] (run_ends = [2, 5])
589        let ree_array = create_run_array_data(vec![2, 5], values);
590
591        let mut mutable = MutableArrayData::new(vec![&ree_array], false, 10);
592
593        // Extend the entire array
594        mutable.try_extend(0, 0, 5).unwrap();
595
596        let result = mutable.freeze();
597
598        // Should have extended correctly
599        assert_eq!(result.len(), 5); // All 5 elements
600
601        // Basic validation that we have the right structure
602        assert!(!result.child_data()[0].is_empty()); // Should have at least one run
603        assert_eq!(result.child_data()[0].len(), result.child_data()[1].len()); // run_ends and values should have same length
604
605        // Should have 2 runs since we have 2 different values
606        assert_eq!(result.child_data()[0].len(), 2);
607        assert_eq!(result.child_data()[1].len(), 2);
608    }
609
610    #[test]
611    fn test_extend_nulls_overflow_i16() {
612        let values = create_int32_array_data(vec![42]);
613        // Start with run end close to max to set up overflow condition
614        let ree_array = create_run_array_data_int16(vec![5], values);
615        let mut mutable = MutableArrayData::new(vec![&ree_array], true, 10);
616
617        // Extend the original data first to initialize state
618        mutable.try_extend(0, 0, 5_usize).unwrap();
619
620        // This should return an error: i16::MAX + 5 > i16::MAX
621        let err = mutable.try_extend_nulls(i16::MAX as usize).unwrap_err();
622        assert!(
623            err.to_string().contains("run end overflow"),
624            "unexpected error: {err}"
625        );
626    }
627
628    #[test]
629    fn test_extend_nulls_overflow_i32() {
630        let values = create_int32_array_data(vec![42]);
631        // Start with run end close to max to set up overflow condition
632        let ree_array = create_run_array_data(vec![10], values);
633        let mut mutable = MutableArrayData::new(vec![&ree_array], true, 10);
634
635        // Extend the original data first to initialize state
636        mutable.try_extend(0, 0, 10_usize).unwrap();
637
638        // This should return an error: (i32::MAX - 10) + 20 > i32::MAX
639        let err = mutable.try_extend_nulls(i32::MAX as usize).unwrap_err();
640        assert!(
641            err.to_string().contains("run end overflow"),
642            "unexpected error: {err}"
643        );
644    }
645
646    #[test]
647    fn test_build_extend_overflow_i16() {
648        // Create a source array with small run that will cause overflow when added
649        let values = create_int32_array_data(vec![10]);
650        let source_array = create_run_array_data_int16(vec![20], values);
651
652        // Create a destination array with run end close to max
653        let dest_values = create_int32_array_data(vec![42]);
654        let dest_array = create_run_array_data_int16(vec![i16::MAX - 5], dest_values);
655
656        let mut mutable = MutableArrayData::new(vec![&source_array, &dest_array], false, 10);
657
658        // First extend the destination array to set up state
659        mutable.try_extend(1, 0, (i16::MAX - 5) as usize).unwrap();
660
661        // This should return an error: (i16::MAX - 5) + 20 > i16::MAX
662        let err = mutable.try_extend(0, 0, 20).unwrap_err();
663        assert!(
664            err.to_string().contains("run end overflow"),
665            "unexpected error: {err}"
666        );
667    }
668
669    #[test]
670    fn test_build_extend_overflow_i32() {
671        // Create a source array with small run that will cause overflow when added
672        let values = create_int32_array_data(vec![10]);
673        let source_array = create_run_array_data(vec![100], values);
674
675        // Create a destination array with run end close to max
676        let dest_values = create_int32_array_data(vec![42]);
677        let dest_array = create_run_array_data(vec![i32::MAX - 50], dest_values);
678
679        let mut mutable = MutableArrayData::new(vec![&source_array, &dest_array], false, 10);
680
681        // First extend the destination array to set up state
682        mutable.try_extend(1, 0, (i32::MAX - 50) as usize).unwrap();
683
684        // This should return an error: (i32::MAX - 50) + 100 > i32::MAX
685        let err = mutable.try_extend(0, 0, 100).unwrap_err();
686        assert!(
687            err.to_string().contains("run end overflow"),
688            "unexpected error: {err}"
689        );
690    }
691}
arrow_data/transform/run.rs

arrow_data/transform/
run.rs
