| 1197 | |
| 1198 | #[test] |
| 1199 | fn test_extend_nulls_dense_union() { |
| 1200 | let fields = UnionFields::try_new( |
| 1201 | vec![0, 1], |
| 1202 | vec![ |
| 1203 | Field::new("i", DataType::Int32, true), |
| 1204 | Field::new("str", DataType::Utf8, true), |
| 1205 | ], |
| 1206 | ) |
| 1207 | .unwrap(); |
| 1208 | |
| 1209 | let type_ids = ScalarBuffer::from(vec![1i8]); |
| 1210 | let offsets = ScalarBuffer::from(vec![0i32]); |
| 1211 | let child_int = Arc::new(Int32Array::new_null(0)) as ArrayRef; |
| 1212 | let child_str = Arc::new(StringArray::from(vec![Some("hello")])) as ArrayRef; |
| 1213 | let union_array = UnionArray::try_new( |
| 1214 | fields.clone(), |
| 1215 | type_ids, |
| 1216 | Some(offsets), |
| 1217 | vec![child_int, child_str], |
| 1218 | ) |
| 1219 | .unwrap(); |
| 1220 | |
| 1221 | let data = union_array.to_data(); |
| 1222 | let mut mutable = MutableArrayData::new(vec![&data], true, 4); |
| 1223 | mutable.extend(0, 0, 1); // copy the first element |
| 1224 | mutable.extend_nulls(2); // add two nulls |
| 1225 | let result = mutable.freeze(); |
| 1226 | |
| 1227 | // Union arrays must not have a null bitmap per Arrow spec |
| 1228 | assert!(result.nulls().is_none()); |
| 1229 | |
| 1230 | let result_array = UnionArray::from(result); |
| 1231 | assert_eq!(result_array.len(), 3); |
| 1232 | // First element is type_id 1 (str) |
| 1233 | assert_eq!(result_array.type_id(0), 1); |
| 1234 | // Null elements use the first type_id (0) |
| 1235 | assert_eq!(result_array.type_id(1), 0); |
| 1236 | assert_eq!(result_array.type_id(2), 0); |
| 1237 | // First child (int) should have 2 null entries from extend_nulls |
| 1238 | assert_eq!(result_array.child(0).len(), 2); |
| 1239 | // Second child (str) should have 1 entry from extend |
| 1240 | assert_eq!(result_array.child(1).len(), 1); |
| 1241 | } |