1use std::sync::Arc;
19
20use arrow_array::{
21 Array, ArrayRef, BinaryArray, BooleanArray, Date32Array, Decimal128Array, FixedSizeBinaryArray,
22 FixedSizeListArray, Float32Array, Float64Array, Int32Array, Int64Array, LargeBinaryArray,
23 LargeListArray, LargeStringArray, ListArray, MapArray, StringArray, StructArray,
24 Time64MicrosecondArray, TimestampMicrosecondArray, TimestampNanosecondArray, new_null_array,
25};
26use arrow_buffer::{BooleanBuffer, NullBuffer};
27use arrow_schema::{DataType, FieldRef, TimeUnit};
28use uuid::Uuid;
29
30use super::get_field_id_from_metadata;
31use crate::spec::{
32 ListType, Literal, Map, MapType, NestedField, PartnerAccessor, PrimitiveLiteral, PrimitiveType,
33 SchemaWithPartnerVisitor, Struct, StructType, Type, VariantType, visit_struct_with_partner,
34 visit_type_with_partner,
35};
36use crate::{Error, ErrorKind, Result};
37
38struct ArrowArrayToIcebergStructConverter;
39
40impl SchemaWithPartnerVisitor<ArrayRef> for ArrowArrayToIcebergStructConverter {
41 type T = Vec<Option<Literal>>;
42
43 fn schema(
44 &mut self,
45 _schema: &crate::spec::Schema,
46 _partner: &ArrayRef,
47 value: Vec<Option<Literal>>,
48 ) -> Result<Vec<Option<Literal>>> {
49 Ok(value)
50 }
51
52 fn field(
53 &mut self,
54 field: &crate::spec::NestedFieldRef,
55 _partner: &ArrayRef,
56 value: Vec<Option<Literal>>,
57 ) -> Result<Vec<Option<Literal>>> {
58 if field.required && value.iter().any(Option::is_none) {
60 return Err(Error::new(
61 ErrorKind::DataInvalid,
62 "The field is required but has null value",
63 )
64 .with_context("field_id", field.id.to_string())
65 .with_context("field_name", &field.name));
66 }
67 Ok(value)
68 }
69
70 fn r#struct(
71 &mut self,
72 _struct: &StructType,
73 array: &ArrayRef,
74 results: Vec<Vec<Option<Literal>>>,
75 ) -> Result<Vec<Option<Literal>>> {
76 let row_len = results.first().map(|column| column.len()).unwrap_or(0);
77 if let Some(col) = results.iter().find(|col| col.len() != row_len) {
78 return Err(Error::new(
79 ErrorKind::DataInvalid,
80 "The struct columns have different row length",
81 )
82 .with_context("first col length", row_len.to_string())
83 .with_context("actual col length", col.len().to_string()));
84 }
85
86 let mut struct_literals = Vec::with_capacity(row_len);
87 let mut columns_iters = results
88 .into_iter()
89 .map(|column| column.into_iter())
90 .collect::<Vec<_>>();
91
92 for i in 0..row_len {
93 let mut literals = Vec::with_capacity(columns_iters.len());
94 for column_iter in columns_iters.iter_mut() {
95 literals.push(column_iter.next().unwrap());
96 }
97 if array.is_null(i) {
98 struct_literals.push(None);
99 } else {
100 struct_literals.push(Some(Literal::Struct(Struct::from_iter(literals))));
101 }
102 }
103
104 Ok(struct_literals)
105 }
106
107 fn list(
108 &mut self,
109 list: &ListType,
110 array: &ArrayRef,
111 elements: Vec<Option<Literal>>,
112 ) -> Result<Vec<Option<Literal>>> {
113 if list.element_field.required && elements.iter().any(Option::is_none) {
114 return Err(Error::new(
115 ErrorKind::DataInvalid,
116 "The list should not have null value",
117 ));
118 }
119 match array.data_type() {
120 DataType::List(_) => {
121 let offset = array
122 .as_any()
123 .downcast_ref::<ListArray>()
124 .ok_or_else(|| {
125 Error::new(ErrorKind::DataInvalid, "The partner is not a list array")
126 })?
127 .offsets();
128 let mut result = Vec::with_capacity(offset.len() - 1);
130 for i in 0..offset.len() - 1 {
131 let start = offset[i] as usize;
132 let end = offset[i + 1] as usize;
133 result.push(Some(Literal::List(elements[start..end].to_vec())));
134 }
135 Ok(result)
136 }
137 DataType::LargeList(_) => {
138 let offset = array
139 .as_any()
140 .downcast_ref::<LargeListArray>()
141 .ok_or_else(|| {
142 Error::new(
143 ErrorKind::DataInvalid,
144 "The partner is not a large list array",
145 )
146 })?
147 .offsets();
148 let mut result = Vec::with_capacity(offset.len() - 1);
150 for i in 0..offset.len() - 1 {
151 let start = offset[i] as usize;
152 let end = offset[i + 1] as usize;
153 result.push(Some(Literal::List(elements[start..end].to_vec())));
154 }
155 Ok(result)
156 }
157 DataType::FixedSizeList(_, len) => {
158 let mut result = Vec::with_capacity(elements.len() / *len as usize);
159 for i in 0..elements.len() / *len as usize {
160 let start = i * *len as usize;
161 let end = (i + 1) * *len as usize;
162 result.push(Some(Literal::List(elements[start..end].to_vec())));
163 }
164 Ok(result)
165 }
166 _ => Err(Error::new(
167 ErrorKind::DataInvalid,
168 "The partner is not a list type",
169 )),
170 }
171 }
172
173 fn map(
174 &mut self,
175 _map: &MapType,
176 partner: &ArrayRef,
177 key_values: Vec<Option<Literal>>,
178 values: Vec<Option<Literal>>,
179 ) -> Result<Vec<Option<Literal>>> {
180 if key_values.len() != values.len() {
182 return Err(Error::new(
183 ErrorKind::DataInvalid,
184 "The key value and value of map should have the same row length",
185 ));
186 }
187
188 let offsets = partner
189 .as_any()
190 .downcast_ref::<MapArray>()
191 .ok_or_else(|| Error::new(ErrorKind::DataInvalid, "The partner is not a map array"))?
192 .offsets();
193 let mut result = Vec::with_capacity(offsets.len() - 1);
195 for i in 0..offsets.len() - 1 {
196 let start = offsets[i] as usize;
197 let end = offsets[i + 1] as usize;
198 let mut map = Map::new();
199 for (key, value) in key_values[start..end].iter().zip(values[start..end].iter()) {
200 map.insert(key.clone().unwrap(), value.clone());
201 }
202 result.push(Some(Literal::Map(map)));
203 }
204 Ok(result)
205 }
206
207 fn primitive(&mut self, p: &PrimitiveType, partner: &ArrayRef) -> Result<Vec<Option<Literal>>> {
208 match p {
209 PrimitiveType::Boolean => {
210 let array = partner
211 .as_any()
212 .downcast_ref::<BooleanArray>()
213 .ok_or_else(|| {
214 Error::new(ErrorKind::DataInvalid, "The partner is not a boolean array")
215 })?;
216 Ok(array.iter().map(|v| v.map(Literal::bool)).collect())
217 }
218 PrimitiveType::Int => {
219 let array = partner
220 .as_any()
221 .downcast_ref::<Int32Array>()
222 .ok_or_else(|| {
223 Error::new(ErrorKind::DataInvalid, "The partner is not a int32 array")
224 })?;
225 Ok(array.iter().map(|v| v.map(Literal::int)).collect())
226 }
227 PrimitiveType::Long => {
228 let array = partner
229 .as_any()
230 .downcast_ref::<Int64Array>()
231 .ok_or_else(|| {
232 Error::new(ErrorKind::DataInvalid, "The partner is not a int64 array")
233 })?;
234 Ok(array.iter().map(|v| v.map(Literal::long)).collect())
235 }
236 PrimitiveType::Float => {
237 let array = partner
238 .as_any()
239 .downcast_ref::<Float32Array>()
240 .ok_or_else(|| {
241 Error::new(ErrorKind::DataInvalid, "The partner is not a float32 array")
242 })?;
243 Ok(array.iter().map(|v| v.map(Literal::float)).collect())
244 }
245 PrimitiveType::Double => {
246 let array = partner
247 .as_any()
248 .downcast_ref::<Float64Array>()
249 .ok_or_else(|| {
250 Error::new(ErrorKind::DataInvalid, "The partner is not a float64 array")
251 })?;
252 Ok(array.iter().map(|v| v.map(Literal::double)).collect())
253 }
254 PrimitiveType::Decimal { precision, scale } => {
255 let array = partner
256 .as_any()
257 .downcast_ref::<Decimal128Array>()
258 .ok_or_else(|| {
259 Error::new(
260 ErrorKind::DataInvalid,
261 "The partner is not a decimal128 array",
262 )
263 })?;
264 if let DataType::Decimal128(arrow_precision, arrow_scale) = array.data_type()
265 && (*arrow_precision as u32 != *precision || *arrow_scale as u32 != *scale)
266 {
267 return Err(Error::new(
268 ErrorKind::DataInvalid,
269 format!(
270 "The precision or scale ({arrow_precision},{arrow_scale}) of arrow decimal128 array is not compatible with iceberg decimal type ({precision},{scale})"
271 ),
272 ));
273 }
274 Ok(array.iter().map(|v| v.map(Literal::decimal)).collect())
275 }
276 PrimitiveType::Date => {
277 let array = partner
278 .as_any()
279 .downcast_ref::<Date32Array>()
280 .ok_or_else(|| {
281 Error::new(ErrorKind::DataInvalid, "The partner is not a date32 array")
282 })?;
283 Ok(array.iter().map(|v| v.map(Literal::date)).collect())
284 }
285 PrimitiveType::Time => {
286 let array = partner
287 .as_any()
288 .downcast_ref::<Time64MicrosecondArray>()
289 .ok_or_else(|| {
290 Error::new(ErrorKind::DataInvalid, "The partner is not a time64 array")
291 })?;
292 Ok(array.iter().map(|v| v.map(Literal::time)).collect())
293 }
294 PrimitiveType::Timestamp => {
295 let array = partner
296 .as_any()
297 .downcast_ref::<TimestampMicrosecondArray>()
298 .ok_or_else(|| {
299 Error::new(
300 ErrorKind::DataInvalid,
301 "The partner is not a timestamp array",
302 )
303 })?;
304 Ok(array.iter().map(|v| v.map(Literal::timestamp)).collect())
305 }
306 PrimitiveType::Timestamptz => {
307 let array = partner
308 .as_any()
309 .downcast_ref::<TimestampMicrosecondArray>()
310 .ok_or_else(|| {
311 Error::new(
312 ErrorKind::DataInvalid,
313 "The partner is not a timestamptz array",
314 )
315 })?;
316 Ok(array.iter().map(|v| v.map(Literal::timestamptz)).collect())
317 }
318 PrimitiveType::TimestampNs => {
319 let array = partner
320 .as_any()
321 .downcast_ref::<TimestampNanosecondArray>()
322 .ok_or_else(|| {
323 Error::new(
324 ErrorKind::DataInvalid,
325 "The partner is not a timestamp_ns array",
326 )
327 })?;
328 Ok(array
329 .iter()
330 .map(|v| v.map(Literal::timestamp_nano))
331 .collect())
332 }
333 PrimitiveType::TimestamptzNs => {
334 let array = partner
335 .as_any()
336 .downcast_ref::<TimestampNanosecondArray>()
337 .ok_or_else(|| {
338 Error::new(
339 ErrorKind::DataInvalid,
340 "The partner is not a timestamptz_ns array",
341 )
342 })?;
343 Ok(array
344 .iter()
345 .map(|v| v.map(Literal::timestamptz_nano))
346 .collect())
347 }
348 PrimitiveType::String => {
349 if let Some(array) = partner.as_any().downcast_ref::<LargeStringArray>() {
350 Ok(array.iter().map(|v| v.map(Literal::string)).collect())
351 } else if let Some(array) = partner.as_any().downcast_ref::<StringArray>() {
352 Ok(array.iter().map(|v| v.map(Literal::string)).collect())
353 } else {
354 Err(Error::new(
355 ErrorKind::DataInvalid,
356 "The partner is not a string array",
357 ))
358 }
359 }
360 PrimitiveType::Uuid => {
361 if let Some(array) = partner.as_any().downcast_ref::<FixedSizeBinaryArray>() {
362 if array.value_length() != 16 {
363 return Err(Error::new(
364 ErrorKind::DataInvalid,
365 "The partner is not a uuid array",
366 ));
367 }
368 Ok(array
369 .iter()
370 .map(|v| {
371 v.map(|v| {
372 Ok(Literal::uuid(Uuid::from_bytes(v.try_into().map_err(
373 |_| {
374 Error::new(
375 ErrorKind::DataInvalid,
376 "Failed to convert binary to uuid",
377 )
378 },
379 )?)))
380 })
381 .transpose()
382 })
383 .collect::<Result<Vec<_>>>()?)
384 } else {
385 Err(Error::new(
386 ErrorKind::DataInvalid,
387 "The partner is not a uuid array",
388 ))
389 }
390 }
391 PrimitiveType::Fixed(len) => {
392 let array = partner
393 .as_any()
394 .downcast_ref::<FixedSizeBinaryArray>()
395 .ok_or_else(|| {
396 Error::new(ErrorKind::DataInvalid, "The partner is not a fixed array")
397 })?;
398 if array.value_length() != *len as i32 {
399 return Err(Error::new(
400 ErrorKind::DataInvalid,
401 "The length of fixed size binary array is not compatible with iceberg fixed type",
402 ));
403 }
404 Ok(array
405 .iter()
406 .map(|v| v.map(|v| Literal::fixed(v.iter().cloned())))
407 .collect())
408 }
409 PrimitiveType::Binary => {
410 if let Some(array) = partner.as_any().downcast_ref::<LargeBinaryArray>() {
411 Ok(array
412 .iter()
413 .map(|v| v.map(|v| Literal::binary(v.to_vec())))
414 .collect())
415 } else if let Some(array) = partner.as_any().downcast_ref::<BinaryArray>() {
416 Ok(array
417 .iter()
418 .map(|v| v.map(|v| Literal::binary(v.to_vec())))
419 .collect())
420 } else {
421 Err(Error::new(
422 ErrorKind::DataInvalid,
423 "The partner is not a binary array",
424 ))
425 }
426 }
427 }
428 }
429
430 fn variant(&mut self, _v: &VariantType, _partner: &ArrayRef) -> Result<Vec<Option<Literal>>> {
431 Err(Error::new(
432 ErrorKind::FeatureUnsupported,
433 "Converting variant Arrow array to Iceberg literal is not supported yet",
434 ))
435 }
436}
437
438#[derive(Clone, Copy, Debug)]
449pub enum FieldMatchMode {
450 Id,
452 Name,
454}
455
456impl FieldMatchMode {
457 pub fn match_field(&self, arrow_field: &FieldRef, iceberg_field: &NestedField) -> bool {
459 match self {
460 FieldMatchMode::Id => get_field_id_from_metadata(arrow_field)
461 .map(|id| id == iceberg_field.id)
462 .unwrap_or(false),
463 FieldMatchMode::Name => arrow_field.name() == &iceberg_field.name,
464 }
465 }
466}
467
468pub struct ArrowArrayAccessor {
470 match_mode: FieldMatchMode,
471}
472
473impl ArrowArrayAccessor {
474 pub fn new() -> Self {
476 Self {
477 match_mode: FieldMatchMode::Id,
478 }
479 }
480
481 pub fn new_with_match_mode(match_mode: FieldMatchMode) -> Self {
483 Self { match_mode }
484 }
485}
486
487impl Default for ArrowArrayAccessor {
488 fn default() -> Self {
489 Self::new()
490 }
491}
492
493impl PartnerAccessor<ArrayRef> for ArrowArrayAccessor {
494 fn struct_partner<'a>(&self, schema_partner: &'a ArrayRef) -> Result<&'a ArrayRef> {
495 if !matches!(schema_partner.data_type(), DataType::Struct(_)) {
496 return Err(Error::new(
497 ErrorKind::DataInvalid,
498 "The schema partner is not a struct type",
499 ));
500 }
501
502 Ok(schema_partner)
503 }
504
505 fn field_partner<'a>(
506 &self,
507 struct_partner: &'a ArrayRef,
508 field: &NestedField,
509 ) -> Result<&'a ArrayRef> {
510 let struct_array = struct_partner
511 .as_any()
512 .downcast_ref::<StructArray>()
513 .ok_or_else(|| {
514 Error::new(
515 ErrorKind::DataInvalid,
516 format!(
517 "The struct partner is not a struct array, partner: {struct_partner:?}"
518 ),
519 )
520 })?;
521
522 let field_pos = struct_array
523 .fields()
524 .iter()
525 .position(|arrow_field| self.match_mode.match_field(arrow_field, field))
526 .ok_or_else(|| {
527 Error::new(
528 ErrorKind::DataInvalid,
529 format!("Field id {} not found in struct array", field.id),
530 )
531 })?;
532
533 Ok(struct_array.column(field_pos))
534 }
535
536 fn list_element_partner<'a>(&self, list_partner: &'a ArrayRef) -> Result<&'a ArrayRef> {
537 match list_partner.data_type() {
538 DataType::List(_) => {
539 let list_array = list_partner
540 .as_any()
541 .downcast_ref::<ListArray>()
542 .ok_or_else(|| {
543 Error::new(
544 ErrorKind::DataInvalid,
545 "The list partner is not a list array",
546 )
547 })?;
548 Ok(list_array.values())
549 }
550 DataType::LargeList(_) => {
551 let list_array = list_partner
552 .as_any()
553 .downcast_ref::<LargeListArray>()
554 .ok_or_else(|| {
555 Error::new(
556 ErrorKind::DataInvalid,
557 "The list partner is not a large list array",
558 )
559 })?;
560 Ok(list_array.values())
561 }
562 DataType::FixedSizeList(_, _) => {
563 let list_array = list_partner
564 .as_any()
565 .downcast_ref::<FixedSizeListArray>()
566 .ok_or_else(|| {
567 Error::new(
568 ErrorKind::DataInvalid,
569 "The list partner is not a fixed size list array",
570 )
571 })?;
572 Ok(list_array.values())
573 }
574 _ => Err(Error::new(
575 ErrorKind::DataInvalid,
576 "The list partner is not a list type",
577 )),
578 }
579 }
580
581 fn map_key_partner<'a>(&self, map_partner: &'a ArrayRef) -> Result<&'a ArrayRef> {
582 let map_array = map_partner
583 .as_any()
584 .downcast_ref::<MapArray>()
585 .ok_or_else(|| {
586 Error::new(ErrorKind::DataInvalid, "The map partner is not a map array")
587 })?;
588 Ok(map_array.keys())
589 }
590
591 fn map_value_partner<'a>(&self, map_partner: &'a ArrayRef) -> Result<&'a ArrayRef> {
592 let map_array = map_partner
593 .as_any()
594 .downcast_ref::<MapArray>()
595 .ok_or_else(|| {
596 Error::new(ErrorKind::DataInvalid, "The map partner is not a map array")
597 })?;
598 Ok(map_array.values())
599 }
600}
601
602pub fn arrow_struct_to_literal(
605 struct_array: &ArrayRef,
606 ty: &StructType,
607) -> Result<Vec<Option<Literal>>> {
608 visit_struct_with_partner(
609 ty,
610 struct_array,
611 &mut ArrowArrayToIcebergStructConverter,
612 &ArrowArrayAccessor::new(),
613 )
614}
615
616pub fn arrow_primitive_to_literal(
619 primitive_array: &ArrayRef,
620 ty: &Type,
621) -> Result<Vec<Option<Literal>>> {
622 visit_type_with_partner(
623 ty,
624 primitive_array,
625 &mut ArrowArrayToIcebergStructConverter,
626 &ArrowArrayAccessor::new(),
627 )
628}
629
630pub(crate) fn create_primitive_array_single_element(
635 data_type: &DataType,
636 prim_lit: Option<&PrimitiveLiteral>,
637) -> Result<ArrayRef> {
638 match (data_type, prim_lit) {
639 (DataType::Boolean, Some(PrimitiveLiteral::Boolean(v))) => {
640 Ok(Arc::new(BooleanArray::from(vec![*v])))
641 }
642 (DataType::Boolean, None) => Ok(Arc::new(BooleanArray::from(vec![Option::<bool>::None]))),
643 (DataType::Int32, Some(PrimitiveLiteral::Int(v))) => {
644 Ok(Arc::new(Int32Array::from(vec![*v])))
645 }
646 (DataType::Int32, None) => Ok(Arc::new(Int32Array::from(vec![Option::<i32>::None]))),
647 (DataType::Date32, Some(PrimitiveLiteral::Int(v))) => {
648 Ok(Arc::new(Date32Array::from(vec![*v])))
649 }
650 (DataType::Date32, None) => Ok(Arc::new(Date32Array::from(vec![Option::<i32>::None]))),
651 (DataType::Int64, Some(PrimitiveLiteral::Long(v))) => {
652 Ok(Arc::new(Int64Array::from(vec![*v])))
653 }
654 (DataType::Int64, None) => Ok(Arc::new(Int64Array::from(vec![Option::<i64>::None]))),
655 (DataType::Timestamp(TimeUnit::Microsecond, timezone), Some(PrimitiveLiteral::Long(v))) => {
656 let array = TimestampMicrosecondArray::from(vec![*v]);
657 if let Some(timezone) = timezone {
658 Ok(Arc::new(array.with_timezone(timezone.clone())))
659 } else {
660 Ok(Arc::new(array))
661 }
662 }
663 (DataType::Timestamp(TimeUnit::Microsecond, timezone), None) => {
664 let array = TimestampMicrosecondArray::from(vec![Option::<i64>::None]);
665 if let Some(timezone) = timezone {
666 Ok(Arc::new(array.with_timezone(timezone.clone())))
667 } else {
668 Ok(Arc::new(array))
669 }
670 }
671 (DataType::Timestamp(TimeUnit::Nanosecond, timezone), Some(PrimitiveLiteral::Long(v))) => {
672 let array = TimestampNanosecondArray::from(vec![*v]);
673 if let Some(timezone) = timezone {
674 Ok(Arc::new(array.with_timezone(timezone.clone())))
675 } else {
676 Ok(Arc::new(array))
677 }
678 }
679 (DataType::Timestamp(TimeUnit::Nanosecond, timezone), None) => {
680 let array = TimestampNanosecondArray::from(vec![Option::<i64>::None]);
681 if let Some(timezone) = timezone {
682 Ok(Arc::new(array.with_timezone(timezone.clone())))
683 } else {
684 Ok(Arc::new(array))
685 }
686 }
687 (DataType::Float32, Some(PrimitiveLiteral::Float(v))) => {
688 Ok(Arc::new(Float32Array::from(vec![v.0])))
689 }
690 (DataType::Float32, None) => Ok(Arc::new(Float32Array::from(vec![Option::<f32>::None]))),
691 (DataType::Float64, Some(PrimitiveLiteral::Double(v))) => {
692 Ok(Arc::new(Float64Array::from(vec![v.0])))
693 }
694 (DataType::Float64, None) => Ok(Arc::new(Float64Array::from(vec![Option::<f64>::None]))),
695 (DataType::Utf8, Some(PrimitiveLiteral::String(v))) => {
696 Ok(Arc::new(StringArray::from(vec![v.as_str()])))
697 }
698 (DataType::Utf8, None) => Ok(Arc::new(StringArray::from(vec![Option::<&str>::None]))),
699 (DataType::Binary, Some(PrimitiveLiteral::Binary(v))) => {
700 Ok(Arc::new(BinaryArray::from_vec(vec![v.as_slice()])))
701 }
702 (DataType::Binary, None) => Ok(Arc::new(BinaryArray::from_opt_vec(vec![
703 Option::<&[u8]>::None,
704 ]))),
705 (DataType::Decimal128(precision, scale), Some(PrimitiveLiteral::Int128(v))) => {
706 let array = Decimal128Array::from(vec![{ *v }])
707 .with_precision_and_scale(*precision, *scale)
708 .map_err(|e| {
709 Error::new(
710 ErrorKind::DataInvalid,
711 format!(
712 "Failed to create Decimal128Array with precision {precision} and scale {scale}: {e}"
713 ),
714 )
715 })?;
716 Ok(Arc::new(array))
717 }
718 (DataType::Decimal128(precision, scale), Some(PrimitiveLiteral::UInt128(v))) => {
719 let array = Decimal128Array::from(vec![*v as i128])
720 .with_precision_and_scale(*precision, *scale)
721 .map_err(|e| {
722 Error::new(
723 ErrorKind::DataInvalid,
724 format!(
725 "Failed to create Decimal128Array with precision {precision} and scale {scale}: {e}"
726 ),
727 )
728 })?;
729 Ok(Arc::new(array))
730 }
731 (DataType::Decimal128(precision, scale), None) => {
732 let array = Decimal128Array::from(vec![Option::<i128>::None])
733 .with_precision_and_scale(*precision, *scale)
734 .map_err(|e| {
735 Error::new(
736 ErrorKind::DataInvalid,
737 format!(
738 "Failed to create Decimal128Array with precision {precision} and scale {scale}: {e}"
739 ),
740 )
741 })?;
742 Ok(Arc::new(array))
743 }
744 (DataType::Struct(fields), None) => {
745 let null_arrays: Vec<ArrayRef> = fields
747 .iter()
748 .map(|f| {
749 match f.data_type() {
752 DataType::Boolean => {
753 Ok(Arc::new(BooleanArray::from(vec![Option::<bool>::None]))
754 as ArrayRef)
755 }
756 DataType::Int32 | DataType::Date32 => {
757 Ok(Arc::new(Int32Array::from(vec![Option::<i32>::None])) as ArrayRef)
758 }
759 DataType::Int64 => {
760 Ok(Arc::new(Int64Array::from(vec![Option::<i64>::None])) as ArrayRef)
761 }
762 DataType::Timestamp(TimeUnit::Microsecond, timezone) => {
763 let array = TimestampMicrosecondArray::from(vec![Option::<i64>::None]);
764 if let Some(timezone) = timezone {
765 Ok(Arc::new(array.with_timezone(timezone.clone())) as ArrayRef)
766 } else {
767 Ok(Arc::new(array) as ArrayRef)
768 }
769 }
770 DataType::Timestamp(TimeUnit::Nanosecond, timezone) => {
771 let array = TimestampNanosecondArray::from(vec![Option::<i64>::None]);
772 if let Some(timezone) = timezone {
773 Ok(Arc::new(array.with_timezone(timezone.clone())) as ArrayRef)
774 } else {
775 Ok(Arc::new(array) as ArrayRef)
776 }
777 }
778 DataType::Float32 => {
779 Ok(Arc::new(Float32Array::from(vec![Option::<f32>::None])) as ArrayRef)
780 }
781 DataType::Float64 => {
782 Ok(Arc::new(Float64Array::from(vec![Option::<f64>::None])) as ArrayRef)
783 }
784 DataType::Utf8 => {
785 Ok(Arc::new(StringArray::from(vec![Option::<&str>::None])) as ArrayRef)
786 }
787 DataType::Binary => {
788 Ok(
789 Arc::new(BinaryArray::from_opt_vec(vec![Option::<&[u8]>::None]))
790 as ArrayRef,
791 )
792 }
793 _ => Err(Error::new(
794 ErrorKind::Unexpected,
795 format!("Unsupported struct field type: {:?}", f.data_type()),
796 )),
797 }
798 })
799 .collect::<Result<Vec<_>>>()?;
800 Ok(Arc::new(StructArray::new(
801 fields.clone(),
802 null_arrays,
803 Some(NullBuffer::new_null(1)),
804 )))
805 }
806 _ => Err(Error::new(
807 ErrorKind::Unexpected,
808 format!("Unsupported constant type combination: {data_type:?} with {prim_lit:?}"),
809 )),
810 }
811}
812
813pub(crate) fn create_primitive_array_repeated(
818 data_type: &DataType,
819 prim_lit: Option<&PrimitiveLiteral>,
820 num_rows: usize,
821) -> Result<ArrayRef> {
822 Ok(match (data_type, prim_lit) {
823 (DataType::Boolean, Some(PrimitiveLiteral::Boolean(value))) => {
825 let buffer = if *value {
826 BooleanBuffer::new_set(num_rows)
827 } else {
828 BooleanBuffer::new_unset(num_rows)
829 };
830 Arc::new(BooleanArray::new(buffer, None))
831 }
832 (DataType::Int32, Some(PrimitiveLiteral::Int(value))) => {
833 Arc::new(Int32Array::from(vec![*value; num_rows]))
834 }
835 (DataType::Date32, Some(PrimitiveLiteral::Int(value))) => {
836 Arc::new(Date32Array::from(vec![*value; num_rows]))
837 }
838 (DataType::Int64, Some(PrimitiveLiteral::Int(value))) => {
839 Arc::new(Int64Array::from(vec![i64::from(*value); num_rows]))
840 }
841 (DataType::Int64, Some(PrimitiveLiteral::Long(value))) => {
842 Arc::new(Int64Array::from(vec![*value; num_rows]))
843 }
844 (
845 DataType::Timestamp(TimeUnit::Microsecond, timezone),
846 Some(PrimitiveLiteral::Long(value)),
847 ) => {
848 let array = TimestampMicrosecondArray::from(vec![*value; num_rows]);
849 if let Some(timezone) = timezone {
850 Arc::new(array.with_timezone(timezone.clone()))
851 } else {
852 Arc::new(array)
853 }
854 }
855 (
856 DataType::Timestamp(TimeUnit::Nanosecond, timezone),
857 Some(PrimitiveLiteral::Long(value)),
858 ) => {
859 let array = TimestampNanosecondArray::from(vec![*value; num_rows]);
860 if let Some(timezone) = timezone {
861 Arc::new(array.with_timezone(timezone.clone()))
862 } else {
863 Arc::new(array)
864 }
865 }
866 (DataType::Float32, Some(PrimitiveLiteral::Float(value))) => {
867 Arc::new(Float32Array::from(vec![value.0; num_rows]))
868 }
869 (DataType::Float64, Some(PrimitiveLiteral::Double(value))) => {
870 Arc::new(Float64Array::from(vec![value.0; num_rows]))
871 }
872 (DataType::Utf8, Some(PrimitiveLiteral::String(value))) => Arc::new(
873 StringArray::from_iter_values(std::iter::repeat_n(value.as_str(), num_rows)),
874 ),
875 (DataType::Binary, Some(PrimitiveLiteral::Binary(value))) => Arc::new(
876 BinaryArray::from_iter_values(std::iter::repeat_n(value.as_slice(), num_rows)),
877 ),
878 (DataType::LargeBinary, Some(PrimitiveLiteral::Binary(value))) => Arc::new(
879 LargeBinaryArray::from_iter_values(std::iter::repeat_n(value.as_slice(), num_rows)),
880 ),
881 (DataType::FixedSizeBinary(len), Some(PrimitiveLiteral::Binary(value))) => {
882 if value.len() != *len as usize {
886 return Err(Error::new(
887 ErrorKind::DataInvalid,
888 format!(
889 "FixedSizeBinary literal length {} does not match declared width {len}",
890 value.len()
891 ),
892 ));
893 }
894
895 Arc::new(
896 FixedSizeBinaryArray::try_from_iter(std::iter::repeat_n(value.as_slice(), num_rows))
897 .map_err(|e| {
898 Error::new(
899 ErrorKind::DataInvalid,
900 format!("Failed to create FixedSizeBinary({len}) array: {e}"),
901 )
902 })?,
903 )
904 }
905 (DataType::Time64(TimeUnit::Microsecond), Some(PrimitiveLiteral::Long(value))) => {
906 Arc::new(Time64MicrosecondArray::from(vec![*value; num_rows]))
907 }
908 (DataType::Decimal128(precision, scale), Some(PrimitiveLiteral::Int128(value))) => {
909 Arc::new(
910 Decimal128Array::from(vec![*value; num_rows])
911 .with_precision_and_scale(*precision, *scale)
912 .map_err(|e| {
913 Error::new(
914 ErrorKind::DataInvalid,
915 format!(
916 "Failed to create Decimal128Array with precision {precision} and scale {scale}: {e}"
917 ),
918 )
919 })?,
920 )
921 }
922 (DataType::Decimal128(precision, scale), Some(PrimitiveLiteral::UInt128(value))) => {
923 Arc::new(
924 Decimal128Array::from(vec![*value as i128; num_rows])
925 .with_precision_and_scale(*precision, *scale)
926 .map_err(|e| {
927 Error::new(
928 ErrorKind::DataInvalid,
929 format!(
930 "Failed to create Decimal128Array with precision {precision} and scale {scale}: {e}"
931 ),
932 )
933 })?,
934 )
935 }
936
937 (DataType::Decimal128(precision, scale), None) => {
939 let vals: Vec<Option<i128>> = vec![None; num_rows];
940 Arc::new(
941 Decimal128Array::from(vals)
942 .with_precision_and_scale(*precision, *scale)
943 .map_err(|e| {
944 Error::new(
945 ErrorKind::DataInvalid,
946 format!(
947 "Failed to create Decimal128Array with precision {precision} and scale {scale}: {e}"
948 ),
949 )
950 })?,
951 )
952 }
953 (DataType::Struct(fields), None) => {
954 let null_arrays: Vec<ArrayRef> = fields
956 .iter()
957 .map(|field| create_primitive_array_repeated(field.data_type(), None, num_rows))
958 .collect::<Result<Vec<_>>>()?;
959
960 Arc::new(StructArray::new(
961 fields.clone(),
962 null_arrays,
963 Some(NullBuffer::new_null(num_rows)),
964 ))
965 }
966 (DataType::Null, _) => Arc::new(arrow_array::NullArray::new(num_rows)),
967
968 (dt, None) => new_null_array(dt, num_rows),
970
971 (dt, _) => {
972 return Err(Error::new(
973 ErrorKind::Unexpected,
974 format!("unexpected target column type {dt}, prim_lit {prim_lit:?}"),
975 ));
976 }
977 })
978}
979
980#[cfg(test)]
981mod test {
982 use std::collections::HashMap;
983 use std::sync::Arc;
984
985 use arrow_array::builder::{Int32Builder, ListBuilder, MapBuilder, StructBuilder};
986 use arrow_array::{
987 ArrayRef, BinaryArray, BooleanArray, Date32Array, Decimal128Array, Float32Array,
988 Float64Array, Int32Array, Int64Array, StringArray, StructArray, Time64MicrosecondArray,
989 TimestampMicrosecondArray, TimestampNanosecondArray,
990 };
991 use arrow_schema::{DataType, Field, Fields, TimeUnit};
992 use parquet::arrow::PARQUET_FIELD_ID_META_KEY;
993
994 use super::*;
995 use crate::spec::{ListType, Literal, MapType, NestedField, PrimitiveType, StructType, Type};
996
997 #[test]
998 fn test_arrow_struct_to_iceberg_struct() {
999 let bool_array = BooleanArray::from(vec![Some(true), Some(false), None]);
1000 let int32_array = Int32Array::from(vec![Some(3), Some(4), None]);
1001 let int64_array = Int64Array::from(vec![Some(5), Some(6), None]);
1002 let float32_array = Float32Array::from(vec![Some(1.1), Some(2.2), None]);
1003 let float64_array = Float64Array::from(vec![Some(3.3), Some(4.4), None]);
1004 let decimal_array = Decimal128Array::from(vec![Some(1000), Some(2000), None])
1005 .with_precision_and_scale(10, 2)
1006 .unwrap();
1007 let date_array = Date32Array::from(vec![Some(18628), Some(18629), None]);
1008 let time_array = Time64MicrosecondArray::from(vec![Some(123456789), Some(987654321), None]);
1009 let timestamp_micro_array = TimestampMicrosecondArray::from(vec![
1010 Some(1622548800000000),
1011 Some(1622635200000000),
1012 None,
1013 ]);
1014 let timestamp_nano_array = TimestampNanosecondArray::from(vec![
1015 Some(1622548800000000000),
1016 Some(1622635200000000000),
1017 None,
1018 ]);
1019 let string_array = StringArray::from(vec![Some("a"), Some("b"), None]);
1020 let binary_array =
1021 BinaryArray::from(vec![Some(b"abc".as_ref()), Some(b"def".as_ref()), None]);
1022
1023 let struct_array = Arc::new(StructArray::from(vec![
1024 (
1025 Arc::new(
1026 Field::new("bool_field", DataType::Boolean, true).with_metadata(HashMap::from(
1027 [(PARQUET_FIELD_ID_META_KEY.to_string(), "0".to_string())],
1028 )),
1029 ),
1030 Arc::new(bool_array) as ArrayRef,
1031 ),
1032 (
1033 Arc::new(
1034 Field::new("int32_field", DataType::Int32, true).with_metadata(HashMap::from(
1035 [(PARQUET_FIELD_ID_META_KEY.to_string(), "2".to_string())],
1036 )),
1037 ),
1038 Arc::new(int32_array) as ArrayRef,
1039 ),
1040 (
1041 Arc::new(
1042 Field::new("int64_field", DataType::Int64, true).with_metadata(HashMap::from(
1043 [(PARQUET_FIELD_ID_META_KEY.to_string(), "3".to_string())],
1044 )),
1045 ),
1046 Arc::new(int64_array) as ArrayRef,
1047 ),
1048 (
1049 Arc::new(
1050 Field::new("float32_field", DataType::Float32, true).with_metadata(
1051 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "4".to_string())]),
1052 ),
1053 ),
1054 Arc::new(float32_array) as ArrayRef,
1055 ),
1056 (
1057 Arc::new(
1058 Field::new("float64_field", DataType::Float64, true).with_metadata(
1059 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "5".to_string())]),
1060 ),
1061 ),
1062 Arc::new(float64_array) as ArrayRef,
1063 ),
1064 (
1065 Arc::new(
1066 Field::new("decimal_field", DataType::Decimal128(10, 2), true).with_metadata(
1067 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "6".to_string())]),
1068 ),
1069 ),
1070 Arc::new(decimal_array) as ArrayRef,
1071 ),
1072 (
1073 Arc::new(
1074 Field::new("date_field", DataType::Date32, true).with_metadata(HashMap::from(
1075 [(PARQUET_FIELD_ID_META_KEY.to_string(), "7".to_string())],
1076 )),
1077 ),
1078 Arc::new(date_array) as ArrayRef,
1079 ),
1080 (
1081 Arc::new(
1082 Field::new("time_field", DataType::Time64(TimeUnit::Microsecond), true)
1083 .with_metadata(HashMap::from([(
1084 PARQUET_FIELD_ID_META_KEY.to_string(),
1085 "8".to_string(),
1086 )])),
1087 ),
1088 Arc::new(time_array) as ArrayRef,
1089 ),
1090 (
1091 Arc::new(
1092 Field::new(
1093 "timestamp_micro_field",
1094 DataType::Timestamp(TimeUnit::Microsecond, None),
1095 true,
1096 )
1097 .with_metadata(HashMap::from([(
1098 PARQUET_FIELD_ID_META_KEY.to_string(),
1099 "9".to_string(),
1100 )])),
1101 ),
1102 Arc::new(timestamp_micro_array) as ArrayRef,
1103 ),
1104 (
1105 Arc::new(
1106 Field::new(
1107 "timestamp_nano_field",
1108 DataType::Timestamp(TimeUnit::Nanosecond, None),
1109 true,
1110 )
1111 .with_metadata(HashMap::from([(
1112 PARQUET_FIELD_ID_META_KEY.to_string(),
1113 "10".to_string(),
1114 )])),
1115 ),
1116 Arc::new(timestamp_nano_array) as ArrayRef,
1117 ),
1118 (
1119 Arc::new(
1120 Field::new("string_field", DataType::Utf8, true).with_metadata(HashMap::from(
1121 [(PARQUET_FIELD_ID_META_KEY.to_string(), "11".to_string())],
1122 )),
1123 ),
1124 Arc::new(string_array) as ArrayRef,
1125 ),
1126 (
1127 Arc::new(
1128 Field::new("binary_field", DataType::Binary, true).with_metadata(
1129 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "12".to_string())]),
1130 ),
1131 ),
1132 Arc::new(binary_array) as ArrayRef,
1133 ),
1134 ])) as ArrayRef;
1135
1136 let iceberg_struct_type = StructType::new(vec![
1137 Arc::new(NestedField::optional(
1138 0,
1139 "bool_field",
1140 Type::Primitive(PrimitiveType::Boolean),
1141 )),
1142 Arc::new(NestedField::optional(
1143 2,
1144 "int32_field",
1145 Type::Primitive(PrimitiveType::Int),
1146 )),
1147 Arc::new(NestedField::optional(
1148 3,
1149 "int64_field",
1150 Type::Primitive(PrimitiveType::Long),
1151 )),
1152 Arc::new(NestedField::optional(
1153 4,
1154 "float32_field",
1155 Type::Primitive(PrimitiveType::Float),
1156 )),
1157 Arc::new(NestedField::optional(
1158 5,
1159 "float64_field",
1160 Type::Primitive(PrimitiveType::Double),
1161 )),
1162 Arc::new(NestedField::optional(
1163 6,
1164 "decimal_field",
1165 Type::Primitive(PrimitiveType::Decimal {
1166 precision: 10,
1167 scale: 2,
1168 }),
1169 )),
1170 Arc::new(NestedField::optional(
1171 7,
1172 "date_field",
1173 Type::Primitive(PrimitiveType::Date),
1174 )),
1175 Arc::new(NestedField::optional(
1176 8,
1177 "time_field",
1178 Type::Primitive(PrimitiveType::Time),
1179 )),
1180 Arc::new(NestedField::optional(
1181 9,
1182 "timestamp_micro_field",
1183 Type::Primitive(PrimitiveType::Timestamp),
1184 )),
1185 Arc::new(NestedField::optional(
1186 10,
1187 "timestamp_nao_field",
1188 Type::Primitive(PrimitiveType::TimestampNs),
1189 )),
1190 Arc::new(NestedField::optional(
1191 11,
1192 "string_field",
1193 Type::Primitive(PrimitiveType::String),
1194 )),
1195 Arc::new(NestedField::optional(
1196 12,
1197 "binary_field",
1198 Type::Primitive(PrimitiveType::Binary),
1199 )),
1200 ]);
1201
1202 let result = arrow_struct_to_literal(&struct_array, &iceberg_struct_type).unwrap();
1203
1204 assert_eq!(result, vec![
1205 Some(Literal::Struct(Struct::from_iter(vec![
1206 Some(Literal::bool(true)),
1207 Some(Literal::int(3)),
1208 Some(Literal::long(5)),
1209 Some(Literal::float(1.1_f32)),
1210 Some(Literal::double(3.3)),
1211 Some(Literal::decimal(1000)),
1212 Some(Literal::date(18628)),
1213 Some(Literal::time(123456789)),
1214 Some(Literal::timestamp(1622548800000000)),
1215 Some(Literal::timestamp_nano(1622548800000000000)),
1216 Some(Literal::string("a".to_string())),
1217 Some(Literal::binary(b"abc".to_vec())),
1218 ]))),
1219 Some(Literal::Struct(Struct::from_iter(vec![
1220 Some(Literal::bool(false)),
1221 Some(Literal::int(4)),
1222 Some(Literal::long(6)),
1223 Some(Literal::float(2.2_f32)),
1224 Some(Literal::double(4.4)),
1225 Some(Literal::decimal(2000)),
1226 Some(Literal::date(18629)),
1227 Some(Literal::time(987654321)),
1228 Some(Literal::timestamp(1622635200000000)),
1229 Some(Literal::timestamp_nano(1622635200000000000)),
1230 Some(Literal::string("b".to_string())),
1231 Some(Literal::binary(b"def".to_vec())),
1232 ]))),
1233 Some(Literal::Struct(Struct::from_iter(vec![
1234 None, None, None, None, None, None, None, None, None, None, None, None,
1235 ]))),
1236 ]);
1237 }
1238
1239 #[test]
1240 fn test_nullable_struct() {
1241 let struct_array = {
1248 let mut builder = StructBuilder::from_fields(
1249 Fields::from(vec![
1250 Field::new("a", DataType::Int32, true).with_metadata(HashMap::from([(
1251 PARQUET_FIELD_ID_META_KEY.to_string(),
1252 "0".to_string(),
1253 )])),
1254 Field::new("b", DataType::Int32, true).with_metadata(HashMap::from([(
1255 PARQUET_FIELD_ID_META_KEY.to_string(),
1256 "1".to_string(),
1257 )])),
1258 ]),
1259 3,
1260 );
1261 builder
1262 .field_builder::<Int32Builder>(0)
1263 .unwrap()
1264 .append_null();
1265 builder
1266 .field_builder::<Int32Builder>(1)
1267 .unwrap()
1268 .append_null();
1269 builder.append(true);
1270
1271 builder
1272 .field_builder::<Int32Builder>(0)
1273 .unwrap()
1274 .append_value(1);
1275 builder
1276 .field_builder::<Int32Builder>(1)
1277 .unwrap()
1278 .append_null();
1279 builder.append(true);
1280
1281 builder
1282 .field_builder::<Int32Builder>(0)
1283 .unwrap()
1284 .append_value(1);
1285 builder
1286 .field_builder::<Int32Builder>(1)
1287 .unwrap()
1288 .append_value(1);
1289 builder.append_null();
1290
1291 Arc::new(builder.finish()) as ArrayRef
1292 };
1293
1294 let iceberg_struct_type = StructType::new(vec![
1295 Arc::new(NestedField::optional(
1296 0,
1297 "a",
1298 Type::Primitive(PrimitiveType::Int),
1299 )),
1300 Arc::new(NestedField::optional(
1301 1,
1302 "b",
1303 Type::Primitive(PrimitiveType::Int),
1304 )),
1305 ]);
1306
1307 let result = arrow_struct_to_literal(&struct_array, &iceberg_struct_type).unwrap();
1308 assert_eq!(result, vec![
1309 Some(Literal::Struct(Struct::from_iter(vec![None, None,]))),
1310 Some(Literal::Struct(Struct::from_iter(vec![
1311 Some(Literal::int(1)),
1312 None,
1313 ]))),
1314 None,
1315 ]);
1316 }
1317
1318 #[test]
1319 fn test_empty_struct() {
1320 let struct_array = Arc::new(StructArray::new_null(Fields::empty(), 3)) as ArrayRef;
1321 let iceberg_struct_type = StructType::new(vec![]);
1322 let result = arrow_struct_to_literal(&struct_array, &iceberg_struct_type).unwrap();
1323 assert_eq!(result, vec![None; 0]);
1324 }
1325
1326 #[test]
1327 fn test_arrow_variant_to_literal_is_unsupported() {
1328 let variant_child = Arc::new(StructArray::from(vec![
1331 (
1332 Arc::new(Field::new("metadata", DataType::Binary, false)),
1333 Arc::new(BinaryArray::from(vec![Some(b"m".as_ref())])) as ArrayRef,
1334 ),
1335 (
1336 Arc::new(Field::new("value", DataType::Binary, false)),
1337 Arc::new(BinaryArray::from(vec![Some(b"v".as_ref())])) as ArrayRef,
1338 ),
1339 ])) as ArrayRef;
1340
1341 let struct_array = Arc::new(StructArray::from(vec![(
1342 Arc::new(
1343 Field::new("v", variant_child.data_type().clone(), false).with_metadata(
1344 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "1".to_string())]),
1345 ),
1346 ),
1347 variant_child,
1348 )])) as ArrayRef;
1349
1350 let ty = StructType::new(vec![
1351 NestedField::required(1, "v", Type::Variant(VariantType)).into(),
1352 ]);
1353
1354 let err = arrow_struct_to_literal(&struct_array, &ty).unwrap_err();
1355 assert_eq!(err.kind(), ErrorKind::FeatureUnsupported);
1356 assert!(
1357 err.to_string()
1358 .contains("Converting variant Arrow array to Iceberg literal is not supported yet"),
1359 "{err}"
1360 );
1361 }
1362
1363 #[test]
1364 fn test_find_field_by_id() {
1365 let field_a_array = Int32Array::from(vec![Some(42), Some(43), None]);
1367 let field_b_array = StringArray::from(vec![Some("value1"), Some("value2"), None]);
1368
1369 let nested_struct_array =
1371 Arc::new(StructArray::from(vec![
1372 (
1373 Arc::new(Field::new("field_a", DataType::Int32, true).with_metadata(
1374 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "1".to_string())]),
1375 )),
1376 Arc::new(field_a_array) as ArrayRef,
1377 ),
1378 (
1379 Arc::new(Field::new("field_b", DataType::Utf8, true).with_metadata(
1380 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "2".to_string())]),
1381 )),
1382 Arc::new(field_b_array) as ArrayRef,
1383 ),
1384 ])) as ArrayRef;
1385
1386 let field_c_array = Int32Array::from(vec![Some(100), Some(200), None]);
1387
1388 let struct_array = Arc::new(StructArray::from(vec![
1390 (
1391 Arc::new(
1392 Field::new(
1393 "nested_struct",
1394 DataType::Struct(Fields::from(vec![
1395 Field::new("field_a", DataType::Int32, true).with_metadata(
1396 HashMap::from([(
1397 PARQUET_FIELD_ID_META_KEY.to_string(),
1398 "1".to_string(),
1399 )]),
1400 ),
1401 Field::new("field_b", DataType::Utf8, true).with_metadata(
1402 HashMap::from([(
1403 PARQUET_FIELD_ID_META_KEY.to_string(),
1404 "2".to_string(),
1405 )]),
1406 ),
1407 ])),
1408 true,
1409 )
1410 .with_metadata(HashMap::from([(
1411 PARQUET_FIELD_ID_META_KEY.to_string(),
1412 "3".to_string(),
1413 )])),
1414 ),
1415 nested_struct_array,
1416 ),
1417 (
1418 Arc::new(Field::new("field_c", DataType::Int32, true).with_metadata(
1419 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "4".to_string())]),
1420 )),
1421 Arc::new(field_c_array) as ArrayRef,
1422 ),
1423 ])) as ArrayRef;
1424
1425 let accessor = ArrowArrayAccessor::new_with_match_mode(FieldMatchMode::Id);
1427
1428 let nested_field = NestedField::optional(
1430 3,
1431 "nested_struct",
1432 Type::Struct(StructType::new(vec![
1433 Arc::new(NestedField::optional(
1434 1,
1435 "field_a",
1436 Type::Primitive(PrimitiveType::Int),
1437 )),
1438 Arc::new(NestedField::optional(
1439 2,
1440 "field_b",
1441 Type::Primitive(PrimitiveType::String),
1442 )),
1443 ])),
1444 );
1445 let nested_partner = accessor
1446 .field_partner(&struct_array, &nested_field)
1447 .unwrap();
1448
1449 let field_a = NestedField::optional(1, "field_a", Type::Primitive(PrimitiveType::Int));
1451 let field_a_partner = accessor.field_partner(nested_partner, &field_a).unwrap();
1452
1453 let int_array = field_a_partner
1455 .as_any()
1456 .downcast_ref::<Int32Array>()
1457 .unwrap();
1458 assert_eq!(int_array.value(0), 42);
1459 assert_eq!(int_array.value(1), 43);
1460 assert!(int_array.is_null(2));
1461 }
1462
1463 #[test]
1464 fn test_find_field_by_name() {
1465 let field_a_array = Int32Array::from(vec![Some(42), Some(43), None]);
1467 let field_b_array = StringArray::from(vec![Some("value1"), Some("value2"), None]);
1468
1469 let nested_struct_array = Arc::new(StructArray::from(vec![
1471 (
1472 Arc::new(Field::new("field_a", DataType::Int32, true)),
1473 Arc::new(field_a_array) as ArrayRef,
1474 ),
1475 (
1476 Arc::new(Field::new("field_b", DataType::Utf8, true)),
1477 Arc::new(field_b_array) as ArrayRef,
1478 ),
1479 ])) as ArrayRef;
1480
1481 let field_c_array = Int32Array::from(vec![Some(100), Some(200), None]);
1482
1483 let struct_array = Arc::new(StructArray::from(vec![
1485 (
1486 Arc::new(Field::new(
1487 "nested_struct",
1488 DataType::Struct(Fields::from(vec![
1489 Field::new("field_a", DataType::Int32, true),
1490 Field::new("field_b", DataType::Utf8, true),
1491 ])),
1492 true,
1493 )),
1494 nested_struct_array,
1495 ),
1496 (
1497 Arc::new(Field::new("field_c", DataType::Int32, true)),
1498 Arc::new(field_c_array) as ArrayRef,
1499 ),
1500 ])) as ArrayRef;
1501
1502 let accessor = ArrowArrayAccessor::new_with_match_mode(FieldMatchMode::Name);
1504
1505 let nested_field = NestedField::optional(
1507 3,
1508 "nested_struct",
1509 Type::Struct(StructType::new(vec![
1510 Arc::new(NestedField::optional(
1511 1,
1512 "field_a",
1513 Type::Primitive(PrimitiveType::Int),
1514 )),
1515 Arc::new(NestedField::optional(
1516 2,
1517 "field_b",
1518 Type::Primitive(PrimitiveType::String),
1519 )),
1520 ])),
1521 );
1522 let nested_partner = accessor
1523 .field_partner(&struct_array, &nested_field)
1524 .unwrap();
1525
1526 let field_a = NestedField::optional(1, "field_a", Type::Primitive(PrimitiveType::Int));
1528 let field_a_partner = accessor.field_partner(nested_partner, &field_a).unwrap();
1529
1530 let int_array = field_a_partner
1532 .as_any()
1533 .downcast_ref::<Int32Array>()
1534 .unwrap();
1535 assert_eq!(int_array.value(0), 42);
1536 assert_eq!(int_array.value(1), 43);
1537 assert!(int_array.is_null(2));
1538 }
1539
1540 #[test]
1541 fn test_complex_nested() {
1542 let struct_type = StructType::new(vec![
1549 Arc::new(NestedField::required(
1550 0,
1551 "A",
1552 Type::List(ListType::new(Arc::new(NestedField::required(
1553 1,
1554 "item",
1555 Type::Struct(StructType::new(vec![
1556 Arc::new(NestedField::required(
1557 2,
1558 "a1",
1559 Type::Primitive(PrimitiveType::Int),
1560 )),
1561 Arc::new(NestedField::required(
1562 3,
1563 "a2",
1564 Type::Primitive(PrimitiveType::Int),
1565 )),
1566 ])),
1567 )))),
1568 )),
1569 Arc::new(NestedField::required(
1570 4,
1571 "B",
1572 Type::List(ListType::new(Arc::new(NestedField::required(
1573 5,
1574 "item",
1575 Type::Map(MapType::new(
1576 NestedField::optional(6, "keys", Type::Primitive(PrimitiveType::Int))
1577 .into(),
1578 NestedField::optional(7, "values", Type::Primitive(PrimitiveType::Int))
1579 .into(),
1580 )),
1581 )))),
1582 )),
1583 Arc::new(NestedField::required(
1584 8,
1585 "C",
1586 Type::List(ListType::new(Arc::new(NestedField::required(
1587 9,
1588 "item",
1589 Type::List(ListType::new(Arc::new(NestedField::optional(
1590 10,
1591 "item",
1592 Type::Primitive(PrimitiveType::Int),
1593 )))),
1594 )))),
1595 )),
1596 ]);
1597
1598 let struct_array =
1604 {
1605 let a_struct_a1_builder = Int32Builder::new();
1606 let a_struct_a2_builder = Int32Builder::new();
1607 let a_struct_builder =
1608 StructBuilder::new(
1609 vec![
1610 Field::new("a1", DataType::Int32, false).with_metadata(HashMap::from(
1611 [(PARQUET_FIELD_ID_META_KEY.to_string(), "2".to_string())],
1612 )),
1613 Field::new("a2", DataType::Int32, false).with_metadata(HashMap::from(
1614 [(PARQUET_FIELD_ID_META_KEY.to_string(), "3".to_string())],
1615 )),
1616 ],
1617 vec![Box::new(a_struct_a1_builder), Box::new(a_struct_a2_builder)],
1618 );
1619 let a_builder = ListBuilder::new(a_struct_builder);
1620
1621 let map_key_builder = Int32Builder::new();
1622 let map_value_builder = Int32Builder::new();
1623 let map_builder = MapBuilder::new(None, map_key_builder, map_value_builder);
1624 let b_builder = ListBuilder::new(map_builder);
1625
1626 let inner_list_item_builder = Int32Builder::new();
1627 let inner_list_builder = ListBuilder::new(inner_list_item_builder);
1628 let c_builder = ListBuilder::new(inner_list_builder);
1629
1630 let mut top_struct_builder = {
1631 let a_struct_type =
1632 DataType::Struct(Fields::from(vec![
1633 Field::new("a1", DataType::Int32, false).with_metadata(HashMap::from(
1634 [(PARQUET_FIELD_ID_META_KEY.to_string(), "2".to_string())],
1635 )),
1636 Field::new("a2", DataType::Int32, false).with_metadata(HashMap::from(
1637 [(PARQUET_FIELD_ID_META_KEY.to_string(), "3".to_string())],
1638 )),
1639 ]));
1640 let a_type =
1641 DataType::List(Arc::new(Field::new("item", a_struct_type.clone(), true)));
1642
1643 let b_map_entry_struct = Field::new(
1644 "entries",
1645 DataType::Struct(Fields::from(vec![
1646 Field::new("keys", DataType::Int32, false),
1647 Field::new("values", DataType::Int32, true),
1648 ])),
1649 false,
1650 );
1651 let b_map_type =
1652 DataType::Map(Arc::new(b_map_entry_struct), false);
1653 let b_type =
1654 DataType::List(Arc::new(Field::new("item", b_map_type.clone(), true)));
1655
1656 let c_inner_list_type =
1657 DataType::List(Arc::new(Field::new("item", DataType::Int32, true)));
1658 let c_type = DataType::List(Arc::new(Field::new(
1659 "item",
1660 c_inner_list_type.clone(),
1661 true,
1662 )));
1663 StructBuilder::new(
1664 Fields::from(vec![
1665 Field::new("A", a_type.clone(), false).with_metadata(HashMap::from([
1666 (PARQUET_FIELD_ID_META_KEY.to_string(), "0".to_string()),
1667 ])),
1668 Field::new("B", b_type.clone(), false).with_metadata(HashMap::from([
1669 (PARQUET_FIELD_ID_META_KEY.to_string(), "4".to_string()),
1670 ])),
1671 Field::new("C", c_type.clone(), false).with_metadata(HashMap::from([
1672 (PARQUET_FIELD_ID_META_KEY.to_string(), "8".to_string()),
1673 ])),
1674 ]),
1675 vec![
1676 Box::new(a_builder),
1677 Box::new(b_builder),
1678 Box::new(c_builder),
1679 ],
1680 )
1681 };
1682
1683 {
1686 let a_builder = top_struct_builder
1687 .field_builder::<ListBuilder<StructBuilder>>(0)
1688 .unwrap();
1689 let struct_builder = a_builder.values();
1690 struct_builder
1691 .field_builder::<Int32Builder>(0)
1692 .unwrap()
1693 .append_value(10);
1694 struct_builder
1695 .field_builder::<Int32Builder>(1)
1696 .unwrap()
1697 .append_value(20);
1698 struct_builder.append(true);
1699 let struct_builder = a_builder.values();
1700 struct_builder
1701 .field_builder::<Int32Builder>(0)
1702 .unwrap()
1703 .append_value(11);
1704 struct_builder
1705 .field_builder::<Int32Builder>(1)
1706 .unwrap()
1707 .append_value(21);
1708 struct_builder.append(true);
1709 a_builder.append(true);
1710 }
1711 {
1712 let b_builder = top_struct_builder
1713 .field_builder::<ListBuilder<MapBuilder<Int32Builder, Int32Builder>>>(1)
1714 .unwrap();
1715 let map_builder = b_builder.values();
1716 map_builder.keys().append_value(1);
1717 map_builder.values().append_value(100);
1718 map_builder.keys().append_value(3);
1719 map_builder.values().append_value(300);
1720 map_builder.append(true).unwrap();
1721
1722 map_builder.keys().append_value(2);
1723 map_builder.values().append_value(200);
1724 map_builder.append(true).unwrap();
1725
1726 b_builder.append(true);
1727 }
1728 {
1729 let c_builder = top_struct_builder
1730 .field_builder::<ListBuilder<ListBuilder<Int32Builder>>>(2)
1731 .unwrap();
1732 let inner_list_builder = c_builder.values();
1733 inner_list_builder.values().append_value(100);
1734 inner_list_builder.values().append_value(101);
1735 inner_list_builder.values().append_value(102);
1736 inner_list_builder.append(true);
1737 let inner_list_builder = c_builder.values();
1738 inner_list_builder.values().append_value(200);
1739 inner_list_builder.values().append_value(201);
1740 inner_list_builder.append(true);
1741 c_builder.append(true);
1742 }
1743 top_struct_builder.append(true);
1744
1745 {
1748 let a_builder = top_struct_builder
1749 .field_builder::<ListBuilder<StructBuilder>>(0)
1750 .unwrap();
1751 let struct_builder = a_builder.values();
1752 struct_builder
1753 .field_builder::<Int32Builder>(0)
1754 .unwrap()
1755 .append_value(12);
1756 struct_builder
1757 .field_builder::<Int32Builder>(1)
1758 .unwrap()
1759 .append_value(22);
1760 struct_builder.append(true);
1761 let struct_builder = a_builder.values();
1762 struct_builder
1763 .field_builder::<Int32Builder>(0)
1764 .unwrap()
1765 .append_value(13);
1766 struct_builder
1767 .field_builder::<Int32Builder>(1)
1768 .unwrap()
1769 .append_value(23);
1770 struct_builder.append(true);
1771 a_builder.append(true);
1772 }
1773 {
1774 let b_builder = top_struct_builder
1775 .field_builder::<ListBuilder<MapBuilder<Int32Builder, Int32Builder>>>(1)
1776 .unwrap();
1777 let map_builder = b_builder.values();
1778 map_builder.keys().append_value(3);
1779 map_builder.values().append_value(300);
1780 map_builder.append(true).unwrap();
1781
1782 b_builder.append(true);
1783 }
1784 {
1785 let c_builder = top_struct_builder
1786 .field_builder::<ListBuilder<ListBuilder<Int32Builder>>>(2)
1787 .unwrap();
1788 let inner_list_builder = c_builder.values();
1789 inner_list_builder.values().append_value(300);
1790 inner_list_builder.values().append_value(301);
1791 inner_list_builder.values().append_value(302);
1792 inner_list_builder.append(true);
1793 let inner_list_builder = c_builder.values();
1794 inner_list_builder.values().append_value(400);
1795 inner_list_builder.values().append_value(401);
1796 inner_list_builder.append(true);
1797 c_builder.append(true);
1798 }
1799 top_struct_builder.append(true);
1800
1801 Arc::new(top_struct_builder.finish()) as ArrayRef
1802 };
1803
1804 let result = arrow_struct_to_literal(&struct_array, &struct_type).unwrap();
1805 assert_eq!(result, vec![
1806 Some(Literal::Struct(Struct::from_iter(vec![
1807 Some(Literal::List(vec![
1808 Some(Literal::Struct(Struct::from_iter(vec![
1809 Some(Literal::int(10)),
1810 Some(Literal::int(20)),
1811 ]))),
1812 Some(Literal::Struct(Struct::from_iter(vec![
1813 Some(Literal::int(11)),
1814 Some(Literal::int(21)),
1815 ]))),
1816 ])),
1817 Some(Literal::List(vec![
1818 Some(Literal::Map(Map::from_iter(vec![
1819 (Literal::int(1), Some(Literal::int(100))),
1820 (Literal::int(3), Some(Literal::int(300))),
1821 ]))),
1822 Some(Literal::Map(Map::from_iter(vec![(
1823 Literal::int(2),
1824 Some(Literal::int(200))
1825 ),]))),
1826 ])),
1827 Some(Literal::List(vec![
1828 Some(Literal::List(vec![
1829 Some(Literal::int(100)),
1830 Some(Literal::int(101)),
1831 Some(Literal::int(102)),
1832 ])),
1833 Some(Literal::List(vec![
1834 Some(Literal::int(200)),
1835 Some(Literal::int(201)),
1836 ])),
1837 ])),
1838 ]))),
1839 Some(Literal::Struct(Struct::from_iter(vec![
1840 Some(Literal::List(vec![
1841 Some(Literal::Struct(Struct::from_iter(vec![
1842 Some(Literal::int(12)),
1843 Some(Literal::int(22)),
1844 ]))),
1845 Some(Literal::Struct(Struct::from_iter(vec![
1846 Some(Literal::int(13)),
1847 Some(Literal::int(23)),
1848 ]))),
1849 ])),
1850 Some(Literal::List(vec![Some(Literal::Map(Map::from_iter(
1851 vec![(Literal::int(3), Some(Literal::int(300))),]
1852 ))),])),
1853 Some(Literal::List(vec![
1854 Some(Literal::List(vec![
1855 Some(Literal::int(300)),
1856 Some(Literal::int(301)),
1857 Some(Literal::int(302)),
1858 ])),
1859 Some(Literal::List(vec![
1860 Some(Literal::int(400)),
1861 Some(Literal::int(401)),
1862 ])),
1863 ])),
1864 ]))),
1865 ]);
1866 }
1867
1868 #[test]
1869 fn test_create_decimal_array_respects_precision() {
1870 let target_precision = 18u8;
1873 let target_scale = 10i8;
1874 let target_type = DataType::Decimal128(target_precision, target_scale);
1875 let value = PrimitiveLiteral::Int128(10000000000);
1876
1877 let array = create_primitive_array_single_element(&target_type, Some(&value))
1878 .expect("Failed to create decimal array");
1879
1880 match array.data_type() {
1881 DataType::Decimal128(precision, scale) => {
1882 assert_eq!(*precision, target_precision);
1883 assert_eq!(*scale, target_scale);
1884 }
1885 other => panic!("Expected Decimal128, got {other:?}"),
1886 }
1887 }
1888
1889 #[test]
1890 fn test_create_decimal_array_repeated_respects_precision() {
1891 let target_precision = 18u8;
1893 let target_scale = 10i8;
1894 let target_type = DataType::Decimal128(target_precision, target_scale);
1895 let value = PrimitiveLiteral::Int128(10000000000);
1896 let num_rows = 5;
1897
1898 let array = create_primitive_array_repeated(&target_type, Some(&value), num_rows)
1899 .expect("Failed to create repeated decimal array");
1900
1901 match array.data_type() {
1902 DataType::Decimal128(precision, scale) => {
1903 assert_eq!(*precision, target_precision);
1904 assert_eq!(*scale, target_scale);
1905 }
1906 other => panic!("Expected Decimal128, got {other:?}"),
1907 }
1908
1909 assert_eq!(array.len(), num_rows);
1910 }
1911
1912 #[test]
1913 fn test_create_timestamp_microsecond_array_repeated() {
1914 let target_type = DataType::Timestamp(TimeUnit::Microsecond, None);
1915 let value = PrimitiveLiteral::Long(1_740_600_000_000_000);
1916 let num_rows = 3;
1917
1918 let array = create_primitive_array_repeated(&target_type, Some(&value), num_rows)
1919 .expect("Failed to create repeated timestamp microsecond array");
1920
1921 assert_eq!(array.data_type(), &target_type);
1922 assert_eq!(array.len(), num_rows);
1923 }
1924
1925 #[test]
1926 fn test_create_timestamp_microsecond_with_timezone_array_repeated() {
1927 let target_type = DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into()));
1928 let value = PrimitiveLiteral::Long(1_740_600_000_000_000);
1929 let num_rows = 2;
1930
1931 let array = create_primitive_array_repeated(&target_type, Some(&value), num_rows)
1932 .expect("Failed to create repeated timestamp microsecond array with timezone");
1933
1934 assert_eq!(array.data_type(), &target_type);
1935 assert_eq!(array.len(), num_rows);
1936 }
1937
1938 #[test]
1939 fn test_create_string_and_binary_arrays_repeated() {
1940 let text = "partition-value-2026";
1941 let bytes: Vec<u8> = vec![0xDE, 0xAD, 0xBE, 0xEF];
1942 let num_rows = 4;
1943
1944 let utf8 = create_primitive_array_repeated(
1945 &DataType::Utf8,
1946 Some(&PrimitiveLiteral::String(text.to_string())),
1947 num_rows,
1948 )
1949 .unwrap();
1950 let utf8 = utf8.as_any().downcast_ref::<StringArray>().unwrap();
1951 assert_eq!(utf8.len(), num_rows);
1952 assert!((0..num_rows).all(|i| utf8.value(i) == text));
1953
1954 let binary = create_primitive_array_repeated(
1955 &DataType::Binary,
1956 Some(&PrimitiveLiteral::Binary(bytes.clone())),
1957 num_rows,
1958 )
1959 .unwrap();
1960 let binary = binary.as_any().downcast_ref::<BinaryArray>().unwrap();
1961 assert_eq!(binary.len(), num_rows);
1962 assert!((0..num_rows).all(|i| binary.value(i) == bytes.as_slice()));
1963
1964 let large = create_primitive_array_repeated(
1965 &DataType::LargeBinary,
1966 Some(&PrimitiveLiteral::Binary(bytes.clone())),
1967 num_rows,
1968 )
1969 .unwrap();
1970 let large = large.as_any().downcast_ref::<LargeBinaryArray>().unwrap();
1971 assert_eq!(large.len(), num_rows);
1972 assert!((0..num_rows).all(|i| large.value(i) == bytes.as_slice()));
1973
1974 let fixed = create_primitive_array_repeated(
1975 &DataType::FixedSizeBinary(bytes.len() as i32),
1976 Some(&PrimitiveLiteral::Binary(bytes.clone())),
1977 num_rows,
1978 )
1979 .unwrap();
1980 let fixed = fixed
1981 .as_any()
1982 .downcast_ref::<FixedSizeBinaryArray>()
1983 .unwrap();
1984 assert_eq!(fixed.len(), num_rows);
1985 assert!((0..num_rows).all(|i| fixed.value(i) == bytes.as_slice()));
1986 }
1987
1988 #[test]
1989 fn test_create_boolean_array_repeated() {
1990 let num_rows = 4;
1991
1992 for value in [true, false] {
1993 let array = create_primitive_array_repeated(
1994 &DataType::Boolean,
1995 Some(&PrimitiveLiteral::Boolean(value)),
1996 num_rows,
1997 )
1998 .unwrap();
1999 let array = array.as_any().downcast_ref::<BooleanArray>().unwrap();
2000 assert_eq!(array.len(), num_rows);
2001 assert_eq!(array.null_count(), 0);
2002 assert!((0..num_rows).all(|i| array.value(i) == value));
2003 }
2004
2005 let empty = create_primitive_array_repeated(
2007 &DataType::Boolean,
2008 Some(&PrimitiveLiteral::Boolean(true)),
2009 0,
2010 )
2011 .unwrap();
2012 assert_eq!(empty.len(), 0);
2013 assert_eq!(empty.null_count(), 0);
2014 }
2015
2016 #[test]
2017 fn test_create_string_array_repeated_empty() {
2018 let array = create_primitive_array_repeated(
2020 &DataType::Utf8,
2021 Some(&PrimitiveLiteral::String("x".to_string())),
2022 0,
2023 )
2024 .unwrap();
2025 assert_eq!(array.len(), 0);
2026 }
2027
2028 #[test]
2029 fn test_create_fixed_size_binary_repeated_wrong_width_errors() {
2030 let err = create_primitive_array_repeated(
2033 &DataType::FixedSizeBinary(4),
2034 Some(&PrimitiveLiteral::Binary(vec![0x01, 0x02, 0x03])),
2035 2,
2036 )
2037 .unwrap_err();
2038 assert!(err.to_string().contains("does not match declared width"));
2039 }
2040
2041 #[test]
2042 fn test_create_fixed_size_binary_repeated_empty_errors() {
2043 let result = create_primitive_array_repeated(
2047 &DataType::FixedSizeBinary(4),
2048 Some(&PrimitiveLiteral::Binary(vec![0xDE, 0xAD, 0xBE, 0xEF])),
2049 0,
2050 );
2051 assert!(result.is_err());
2052 }
2053}