Fix: Fix null handling in CometVector implementations #2643
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -16,11 +16,10 @@ | |
| // under the License. | ||
|
|
||
| use arrow::array::{make_array, Array, ArrayRef, GenericListArray, Int32Array, OffsetSizeTrait}; | ||
| use arrow::datatypes::{DataType, Schema}; | ||
| use arrow::{ | ||
| array::{as_primitive_array, Capacities, MutableArrayData}, | ||
| buffer::{NullBuffer, OffsetBuffer}, | ||
| record_batch::RecordBatch, | ||
| }; | ||
| use datafusion::common::{ | ||
|
|
@@ -198,114 +197,124 @@ fn array_insert<O: OffsetSizeTrait>( | |
| pos_array: &ArrayRef, | ||
| legacy_mode: bool, | ||
| ) -> DataFusionResult<ColumnarValue> { | ||
| // Implementation aligned with Arrow's half-open offset ranges and Spark semantics. | ||
|
|
||
| let values = list_array.values(); | ||
| let offsets = list_array.offsets(); | ||
| let values_data = values.to_data(); | ||
| let item_data = items_array.to_data(); | ||
|
|
||
| // Estimate capacity (original values + inserted items upper bound) | ||
| let new_capacity = Capacities::Array(values_data.len() + item_data.len()); | ||
|
|
||
| let mut mutable_values = | ||
| MutableArrayData::with_capacities(vec![&values_data, &item_data], true, new_capacity); | ||
|
|
||
| // New offsets and top-level list validity bitmap | ||
| let mut new_offsets = Vec::with_capacity(list_array.len() + 1); | ||
| new_offsets.push(O::usize_as(0)); | ||
| let mut list_valid = Vec::<bool>::with_capacity(list_array.len()); | ||
|
|
||
| // Spark supports only Int32 position indices | ||
| let pos_data: &Int32Array = as_primitive_array(&pos_array); | ||
|
|
||
| for (row_index, window) in offsets.windows(2).enumerate() { | ||
| let start = window[0].as_usize(); | ||
| let end = window[1].as_usize(); | ||
| let len = end - start; | ||
| let pos = pos_data.value(row_index); | ||
|
Member
There was a problem hiding this comment. Should |
||
|
|
||
| if list_array.is_null(row_index) { | ||
| // Top-level list row is NULL: do not write any child values and do not advance offset | ||
| new_offsets.push(new_offsets[row_index]); | ||
| list_valid.push(false); | ||
| continue; | ||
| } | ||
|
|
||
| if pos == 0 { | ||
| return Err(DataFusionError::Internal( | ||
| "Position for array_insert should be greater or less than zero".to_string(), | ||
| )); | ||
| } | ||
|
|
||
| let final_len: usize; | ||
|
|
||
| if pos > 0 { | ||
| // Positive index (1-based) | ||
| let pos1 = pos as usize; | ||
| if pos1 <= len + 1 { | ||
| // In-range insertion (including appending to end) | ||
| let corrected = pos1 - 1; // 0-based insertion point | ||
| mutable_values.extend(0, start, start + corrected); | ||
| mutable_values.extend(1, row_index, row_index + 1); | ||
| mutable_values.extend(0, start + corrected, end); | ||
| final_len = len + 1; | ||
| } else { | ||
| // Beyond end: pad with nulls then insert | ||
| let corrected = pos1 - 1; | ||
| let padding = corrected - len; | ||
| mutable_values.extend(0, start, end); | ||
| mutable_values.extend_nulls(padding); | ||
| mutable_values.extend(1, row_index, row_index + 1); | ||
| final_len = corrected + 1; // equals pos1 | ||
| } | ||
| } else { | ||
| // Negative index (1-based from the end) | ||
| let k = (-pos) as usize; | ||
|
|
||
| if k <= len { | ||
| // In-range negative insertion | ||
| // Non-legacy: -1 behaves like append to end (corrected = len - k + 1) | ||
| // Legacy: -1 behaves like insert before the last element (corrected = len - k) | ||
| let base_offset = if legacy_mode { 0 } else { 1 }; | ||
| let corrected = len - k + base_offset; | ||
| mutable_values.extend(0, start, start + corrected); | ||
| mutable_values.extend(1, row_index, row_index + 1); | ||
| mutable_values.extend(0, start + corrected, end); | ||
| final_len = len + 1; | ||
| } else { | ||
| // Negative index beyond the start (Spark-specific behavior): | ||
| // Place item first, then pad with nulls, then append the original array. | ||
| // Final length = k + base_offset, where base_offset = 1 in legacy mode, otherwise 0. | ||
| let base_offset = if legacy_mode { 1 } else { 0 }; | ||
| let target_len = k + base_offset; | ||
| let padding = target_len.saturating_sub(len + 1); | ||
| mutable_values.extend(1, row_index, row_index + 1); // insert item first | ||
| mutable_values.extend_nulls(padding); // pad nulls | ||
| mutable_values.extend(0, start, end); // append original values | ||
| final_len = target_len; | ||
| } | ||
| } | ||
|
|
||
| if final_len > MAX_ROUNDED_ARRAY_LENGTH { | ||
| return Err(DataFusionError::Internal(format!( | ||
| "Max array length in Spark is {MAX_ROUNDED_ARRAY_LENGTH}, but got {final_len}" | ||
| ))); | ||
| } | ||
|
|
||
| let prev = new_offsets[row_index].as_usize(); | ||
| new_offsets.push(O::usize_as(prev + final_len)); | ||
| list_valid.push(true); | ||
| } | ||
|
|
||
| let child = make_array(mutable_values.freeze()); | ||
|
|
||
| // Reuse the original list element field (name/type/nullability) | ||
| let elem_field = match list_array.data_type() { | ||
| DataType::List(field) => Arc::clone(field), | ||
| DataType::LargeList(field) => Arc::clone(field), | ||
| _ => unreachable!(), | ||
| }; | ||
|
|
||
| // Build the resulting list array | ||
| let new_list = GenericListArray::<O>::try_new( | ||
| elem_field, | ||
| OffsetBuffer::new(new_offsets.into()), | ||
| child, | ||
| Some(NullBuffer::new(list_valid.into())), | ||
| )?; | ||
|
|
||
| Ok(ColumnarValue::Array(Arc::new(new_list))) | ||
| } | ||
|
|
||
| impl Display for ArrayInsert { | ||
|
|
@@ -442,4 +451,37 @@ mod test { | |
|
|
||
| Ok(()) | ||
| } | ||
|
|
||
| #[test] | ||
| fn test_array_insert_bug_repro_null_item_pos1_fixed() -> Result<()> { | ||
| use arrow::array::{Array, ArrayRef, Int32Array, ListArray}; | ||
| use arrow::datatypes::Int32Type; | ||
|
|
||
| // row0 = [0, null, 0] | ||
| // row1 = [1, null, 1] | ||
| let list = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![ | ||
| Some(vec![Some(0), None, Some(0)]), | ||
| Some(vec![Some(1), None, Some(1)]), | ||
| ]); | ||
|
|
||
| let positions = Int32Array::from(vec![1, 1]); | ||
| let items = Int32Array::from(vec![None, None]); | ||
|
|
||
| let ColumnarValue::Array(result) = array_insert( | ||
| &list, | ||
| &(Arc::new(items) as ArrayRef), | ||
| &(Arc::new(positions) as ArrayRef), | ||
| false, // legacy_mode = false | ||
| )? | ||
| else { | ||
| unreachable!() | ||
| }; | ||
|
|
||
| let expected = ListArray::from_iter_primitive::<Int32Type, _, _>(vec![ | ||
| Some(vec![None, Some(0), None, Some(0)]), | ||
| Some(vec![None, Some(1), None, Some(1)]), | ||
| ]); | ||
| assert_eq!(&result.to_data(), &expected.to_data()); | ||
| Ok(()) | ||
| } | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The version fixed by ChatGPT :)