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,
25};
26use arrow_buffer::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))) => {
824 Arc::new(BooleanArray::from(vec![*value; num_rows]))
825 }
826 (DataType::Boolean, None) => {
827 let vals: Vec<Option<bool>> = vec![None; num_rows];
828 Arc::new(BooleanArray::from(vals))
829 }
830 (DataType::Int32, Some(PrimitiveLiteral::Int(value))) => {
831 Arc::new(Int32Array::from(vec![*value; num_rows]))
832 }
833 (DataType::Int32, None) => {
834 let vals: Vec<Option<i32>> = vec![None; num_rows];
835 Arc::new(Int32Array::from(vals))
836 }
837 (DataType::Date32, Some(PrimitiveLiteral::Int(value))) => {
838 Arc::new(Date32Array::from(vec![*value; num_rows]))
839 }
840 (DataType::Date32, None) => {
841 let vals: Vec<Option<i32>> = vec![None; num_rows];
842 Arc::new(Date32Array::from(vals))
843 }
844 (DataType::Int64, Some(PrimitiveLiteral::Int(value))) => {
845 Arc::new(Int64Array::from(vec![i64::from(*value); num_rows]))
846 }
847 (DataType::Int64, Some(PrimitiveLiteral::Long(value))) => {
848 Arc::new(Int64Array::from(vec![*value; num_rows]))
849 }
850 (DataType::Int64, None) => {
851 let vals: Vec<Option<i64>> = vec![None; num_rows];
852 Arc::new(Int64Array::from(vals))
853 }
854 (
855 DataType::Timestamp(TimeUnit::Microsecond, timezone),
856 Some(PrimitiveLiteral::Long(value)),
857 ) => {
858 let array = TimestampMicrosecondArray::from(vec![*value; num_rows]);
859 if let Some(timezone) = timezone {
860 Arc::new(array.with_timezone(timezone.clone()))
861 } else {
862 Arc::new(array)
863 }
864 }
865 (DataType::Timestamp(TimeUnit::Microsecond, timezone), None) => {
866 let vals: Vec<Option<i64>> = vec![None; num_rows];
867 let array = TimestampMicrosecondArray::from(vals);
868 if let Some(timezone) = timezone {
869 Arc::new(array.with_timezone(timezone.clone()))
870 } else {
871 Arc::new(array)
872 }
873 }
874 (
875 DataType::Timestamp(TimeUnit::Nanosecond, timezone),
876 Some(PrimitiveLiteral::Long(value)),
877 ) => {
878 let array = TimestampNanosecondArray::from(vec![*value; num_rows]);
879 if let Some(timezone) = timezone {
880 Arc::new(array.with_timezone(timezone.clone()))
881 } else {
882 Arc::new(array)
883 }
884 }
885 (DataType::Timestamp(TimeUnit::Nanosecond, timezone), None) => {
886 let vals: Vec<Option<i64>> = vec![None; num_rows];
887 let array = TimestampNanosecondArray::from(vals);
888 if let Some(timezone) = timezone {
889 Arc::new(array.with_timezone(timezone.clone()))
890 } else {
891 Arc::new(array)
892 }
893 }
894 (DataType::Float32, Some(PrimitiveLiteral::Float(value))) => {
895 Arc::new(Float32Array::from(vec![value.0; num_rows]))
896 }
897 (DataType::Float32, None) => {
898 let vals: Vec<Option<f32>> = vec![None; num_rows];
899 Arc::new(Float32Array::from(vals))
900 }
901 (DataType::Float64, Some(PrimitiveLiteral::Double(value))) => {
902 Arc::new(Float64Array::from(vec![value.0; num_rows]))
903 }
904 (DataType::Float64, None) => {
905 let vals: Vec<Option<f64>> = vec![None; num_rows];
906 Arc::new(Float64Array::from(vals))
907 }
908 (DataType::Utf8, Some(PrimitiveLiteral::String(value))) => {
909 Arc::new(StringArray::from(vec![value.clone(); num_rows]))
910 }
911 (DataType::Utf8, None) => {
912 let vals: Vec<Option<String>> = vec![None; num_rows];
913 Arc::new(StringArray::from(vals))
914 }
915 (DataType::Binary, Some(PrimitiveLiteral::Binary(value))) => {
916 Arc::new(BinaryArray::from_vec(vec![value; num_rows]))
917 }
918 (DataType::Binary, None) => {
919 let vals: Vec<Option<&[u8]>> = vec![None; num_rows];
920 Arc::new(BinaryArray::from_opt_vec(vals))
921 }
922 (DataType::Decimal128(precision, scale), Some(PrimitiveLiteral::Int128(value))) => {
923 Arc::new(
924 Decimal128Array::from(vec![*value; 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 (DataType::Decimal128(precision, scale), Some(PrimitiveLiteral::UInt128(value))) => {
937 Arc::new(
938 Decimal128Array::from(vec![*value as i128; num_rows])
939 .with_precision_and_scale(*precision, *scale)
940 .map_err(|e| {
941 Error::new(
942 ErrorKind::DataInvalid,
943 format!(
944 "Failed to create Decimal128Array with precision {precision} and scale {scale}: {e}"
945 ),
946 )
947 })?,
948 )
949 }
950 (DataType::Decimal128(precision, scale), None) => {
951 let vals: Vec<Option<i128>> = vec![None; num_rows];
952 Arc::new(
953 Decimal128Array::from(vals)
954 .with_precision_and_scale(*precision, *scale)
955 .map_err(|e| {
956 Error::new(
957 ErrorKind::DataInvalid,
958 format!(
959 "Failed to create Decimal128Array with precision {precision} and scale {scale}: {e}"
960 ),
961 )
962 })?,
963 )
964 }
965 (DataType::Struct(fields), None) => {
966 let null_arrays: Vec<ArrayRef> = fields
968 .iter()
969 .map(|field| create_primitive_array_repeated(field.data_type(), &None, num_rows))
970 .collect::<Result<Vec<_>>>()?;
971
972 Arc::new(StructArray::new(
973 fields.clone(),
974 null_arrays,
975 Some(NullBuffer::new_null(num_rows)),
976 ))
977 }
978 (DataType::Null, _) => Arc::new(arrow_array::NullArray::new(num_rows)),
979 (dt, _) => {
980 return Err(Error::new(
981 ErrorKind::Unexpected,
982 format!("unexpected target column type {dt}, prim_lit {prim_lit:?}"),
983 ));
984 }
985 })
986}
987
988#[cfg(test)]
989mod test {
990 use std::collections::HashMap;
991 use std::sync::Arc;
992
993 use arrow_array::builder::{Int32Builder, ListBuilder, MapBuilder, StructBuilder};
994 use arrow_array::{
995 ArrayRef, BinaryArray, BooleanArray, Date32Array, Decimal128Array, Float32Array,
996 Float64Array, Int32Array, Int64Array, StringArray, StructArray, Time64MicrosecondArray,
997 TimestampMicrosecondArray, TimestampNanosecondArray,
998 };
999 use arrow_schema::{DataType, Field, Fields, TimeUnit};
1000 use parquet::arrow::PARQUET_FIELD_ID_META_KEY;
1001
1002 use super::*;
1003 use crate::spec::{ListType, Literal, MapType, NestedField, PrimitiveType, StructType, Type};
1004
1005 #[test]
1006 fn test_arrow_struct_to_iceberg_struct() {
1007 let bool_array = BooleanArray::from(vec![Some(true), Some(false), None]);
1008 let int32_array = Int32Array::from(vec![Some(3), Some(4), None]);
1009 let int64_array = Int64Array::from(vec![Some(5), Some(6), None]);
1010 let float32_array = Float32Array::from(vec![Some(1.1), Some(2.2), None]);
1011 let float64_array = Float64Array::from(vec![Some(3.3), Some(4.4), None]);
1012 let decimal_array = Decimal128Array::from(vec![Some(1000), Some(2000), None])
1013 .with_precision_and_scale(10, 2)
1014 .unwrap();
1015 let date_array = Date32Array::from(vec![Some(18628), Some(18629), None]);
1016 let time_array = Time64MicrosecondArray::from(vec![Some(123456789), Some(987654321), None]);
1017 let timestamp_micro_array = TimestampMicrosecondArray::from(vec![
1018 Some(1622548800000000),
1019 Some(1622635200000000),
1020 None,
1021 ]);
1022 let timestamp_nano_array = TimestampNanosecondArray::from(vec![
1023 Some(1622548800000000000),
1024 Some(1622635200000000000),
1025 None,
1026 ]);
1027 let string_array = StringArray::from(vec![Some("a"), Some("b"), None]);
1028 let binary_array =
1029 BinaryArray::from(vec![Some(b"abc".as_ref()), Some(b"def".as_ref()), None]);
1030
1031 let struct_array = Arc::new(StructArray::from(vec![
1032 (
1033 Arc::new(
1034 Field::new("bool_field", DataType::Boolean, true).with_metadata(HashMap::from(
1035 [(PARQUET_FIELD_ID_META_KEY.to_string(), "0".to_string())],
1036 )),
1037 ),
1038 Arc::new(bool_array) as ArrayRef,
1039 ),
1040 (
1041 Arc::new(
1042 Field::new("int32_field", DataType::Int32, true).with_metadata(HashMap::from(
1043 [(PARQUET_FIELD_ID_META_KEY.to_string(), "2".to_string())],
1044 )),
1045 ),
1046 Arc::new(int32_array) as ArrayRef,
1047 ),
1048 (
1049 Arc::new(
1050 Field::new("int64_field", DataType::Int64, true).with_metadata(HashMap::from(
1051 [(PARQUET_FIELD_ID_META_KEY.to_string(), "3".to_string())],
1052 )),
1053 ),
1054 Arc::new(int64_array) as ArrayRef,
1055 ),
1056 (
1057 Arc::new(
1058 Field::new("float32_field", DataType::Float32, true).with_metadata(
1059 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "4".to_string())]),
1060 ),
1061 ),
1062 Arc::new(float32_array) as ArrayRef,
1063 ),
1064 (
1065 Arc::new(
1066 Field::new("float64_field", DataType::Float64, true).with_metadata(
1067 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "5".to_string())]),
1068 ),
1069 ),
1070 Arc::new(float64_array) as ArrayRef,
1071 ),
1072 (
1073 Arc::new(
1074 Field::new("decimal_field", DataType::Decimal128(10, 2), true).with_metadata(
1075 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "6".to_string())]),
1076 ),
1077 ),
1078 Arc::new(decimal_array) as ArrayRef,
1079 ),
1080 (
1081 Arc::new(
1082 Field::new("date_field", DataType::Date32, true).with_metadata(HashMap::from(
1083 [(PARQUET_FIELD_ID_META_KEY.to_string(), "7".to_string())],
1084 )),
1085 ),
1086 Arc::new(date_array) as ArrayRef,
1087 ),
1088 (
1089 Arc::new(
1090 Field::new("time_field", DataType::Time64(TimeUnit::Microsecond), true)
1091 .with_metadata(HashMap::from([(
1092 PARQUET_FIELD_ID_META_KEY.to_string(),
1093 "8".to_string(),
1094 )])),
1095 ),
1096 Arc::new(time_array) as ArrayRef,
1097 ),
1098 (
1099 Arc::new(
1100 Field::new(
1101 "timestamp_micro_field",
1102 DataType::Timestamp(TimeUnit::Microsecond, None),
1103 true,
1104 )
1105 .with_metadata(HashMap::from([(
1106 PARQUET_FIELD_ID_META_KEY.to_string(),
1107 "9".to_string(),
1108 )])),
1109 ),
1110 Arc::new(timestamp_micro_array) as ArrayRef,
1111 ),
1112 (
1113 Arc::new(
1114 Field::new(
1115 "timestamp_nano_field",
1116 DataType::Timestamp(TimeUnit::Nanosecond, None),
1117 true,
1118 )
1119 .with_metadata(HashMap::from([(
1120 PARQUET_FIELD_ID_META_KEY.to_string(),
1121 "10".to_string(),
1122 )])),
1123 ),
1124 Arc::new(timestamp_nano_array) as ArrayRef,
1125 ),
1126 (
1127 Arc::new(
1128 Field::new("string_field", DataType::Utf8, true).with_metadata(HashMap::from(
1129 [(PARQUET_FIELD_ID_META_KEY.to_string(), "11".to_string())],
1130 )),
1131 ),
1132 Arc::new(string_array) as ArrayRef,
1133 ),
1134 (
1135 Arc::new(
1136 Field::new("binary_field", DataType::Binary, true).with_metadata(
1137 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "12".to_string())]),
1138 ),
1139 ),
1140 Arc::new(binary_array) as ArrayRef,
1141 ),
1142 ])) as ArrayRef;
1143
1144 let iceberg_struct_type = StructType::new(vec![
1145 Arc::new(NestedField::optional(
1146 0,
1147 "bool_field",
1148 Type::Primitive(PrimitiveType::Boolean),
1149 )),
1150 Arc::new(NestedField::optional(
1151 2,
1152 "int32_field",
1153 Type::Primitive(PrimitiveType::Int),
1154 )),
1155 Arc::new(NestedField::optional(
1156 3,
1157 "int64_field",
1158 Type::Primitive(PrimitiveType::Long),
1159 )),
1160 Arc::new(NestedField::optional(
1161 4,
1162 "float32_field",
1163 Type::Primitive(PrimitiveType::Float),
1164 )),
1165 Arc::new(NestedField::optional(
1166 5,
1167 "float64_field",
1168 Type::Primitive(PrimitiveType::Double),
1169 )),
1170 Arc::new(NestedField::optional(
1171 6,
1172 "decimal_field",
1173 Type::Primitive(PrimitiveType::Decimal {
1174 precision: 10,
1175 scale: 2,
1176 }),
1177 )),
1178 Arc::new(NestedField::optional(
1179 7,
1180 "date_field",
1181 Type::Primitive(PrimitiveType::Date),
1182 )),
1183 Arc::new(NestedField::optional(
1184 8,
1185 "time_field",
1186 Type::Primitive(PrimitiveType::Time),
1187 )),
1188 Arc::new(NestedField::optional(
1189 9,
1190 "timestamp_micro_field",
1191 Type::Primitive(PrimitiveType::Timestamp),
1192 )),
1193 Arc::new(NestedField::optional(
1194 10,
1195 "timestamp_nao_field",
1196 Type::Primitive(PrimitiveType::TimestampNs),
1197 )),
1198 Arc::new(NestedField::optional(
1199 11,
1200 "string_field",
1201 Type::Primitive(PrimitiveType::String),
1202 )),
1203 Arc::new(NestedField::optional(
1204 12,
1205 "binary_field",
1206 Type::Primitive(PrimitiveType::Binary),
1207 )),
1208 ]);
1209
1210 let result = arrow_struct_to_literal(&struct_array, &iceberg_struct_type).unwrap();
1211
1212 assert_eq!(result, vec![
1213 Some(Literal::Struct(Struct::from_iter(vec![
1214 Some(Literal::bool(true)),
1215 Some(Literal::int(3)),
1216 Some(Literal::long(5)),
1217 Some(Literal::float(1.1)),
1218 Some(Literal::double(3.3)),
1219 Some(Literal::decimal(1000)),
1220 Some(Literal::date(18628)),
1221 Some(Literal::time(123456789)),
1222 Some(Literal::timestamp(1622548800000000)),
1223 Some(Literal::timestamp_nano(1622548800000000000)),
1224 Some(Literal::string("a".to_string())),
1225 Some(Literal::binary(b"abc".to_vec())),
1226 ]))),
1227 Some(Literal::Struct(Struct::from_iter(vec![
1228 Some(Literal::bool(false)),
1229 Some(Literal::int(4)),
1230 Some(Literal::long(6)),
1231 Some(Literal::float(2.2)),
1232 Some(Literal::double(4.4)),
1233 Some(Literal::decimal(2000)),
1234 Some(Literal::date(18629)),
1235 Some(Literal::time(987654321)),
1236 Some(Literal::timestamp(1622635200000000)),
1237 Some(Literal::timestamp_nano(1622635200000000000)),
1238 Some(Literal::string("b".to_string())),
1239 Some(Literal::binary(b"def".to_vec())),
1240 ]))),
1241 Some(Literal::Struct(Struct::from_iter(vec![
1242 None, None, None, None, None, None, None, None, None, None, None, None,
1243 ]))),
1244 ]);
1245 }
1246
1247 #[test]
1248 fn test_nullable_struct() {
1249 let struct_array = {
1256 let mut builder = StructBuilder::from_fields(
1257 Fields::from(vec![
1258 Field::new("a", DataType::Int32, true).with_metadata(HashMap::from([(
1259 PARQUET_FIELD_ID_META_KEY.to_string(),
1260 "0".to_string(),
1261 )])),
1262 Field::new("b", DataType::Int32, true).with_metadata(HashMap::from([(
1263 PARQUET_FIELD_ID_META_KEY.to_string(),
1264 "1".to_string(),
1265 )])),
1266 ]),
1267 3,
1268 );
1269 builder
1270 .field_builder::<Int32Builder>(0)
1271 .unwrap()
1272 .append_null();
1273 builder
1274 .field_builder::<Int32Builder>(1)
1275 .unwrap()
1276 .append_null();
1277 builder.append(true);
1278
1279 builder
1280 .field_builder::<Int32Builder>(0)
1281 .unwrap()
1282 .append_value(1);
1283 builder
1284 .field_builder::<Int32Builder>(1)
1285 .unwrap()
1286 .append_null();
1287 builder.append(true);
1288
1289 builder
1290 .field_builder::<Int32Builder>(0)
1291 .unwrap()
1292 .append_value(1);
1293 builder
1294 .field_builder::<Int32Builder>(1)
1295 .unwrap()
1296 .append_value(1);
1297 builder.append_null();
1298
1299 Arc::new(builder.finish()) as ArrayRef
1300 };
1301
1302 let iceberg_struct_type = StructType::new(vec![
1303 Arc::new(NestedField::optional(
1304 0,
1305 "a",
1306 Type::Primitive(PrimitiveType::Int),
1307 )),
1308 Arc::new(NestedField::optional(
1309 1,
1310 "b",
1311 Type::Primitive(PrimitiveType::Int),
1312 )),
1313 ]);
1314
1315 let result = arrow_struct_to_literal(&struct_array, &iceberg_struct_type).unwrap();
1316 assert_eq!(result, vec![
1317 Some(Literal::Struct(Struct::from_iter(vec![None, None,]))),
1318 Some(Literal::Struct(Struct::from_iter(vec![
1319 Some(Literal::int(1)),
1320 None,
1321 ]))),
1322 None,
1323 ]);
1324 }
1325
1326 #[test]
1327 fn test_empty_struct() {
1328 let struct_array = Arc::new(StructArray::new_null(Fields::empty(), 3)) as ArrayRef;
1329 let iceberg_struct_type = StructType::new(vec![]);
1330 let result = arrow_struct_to_literal(&struct_array, &iceberg_struct_type).unwrap();
1331 assert_eq!(result, vec![None; 0]);
1332 }
1333
1334 #[test]
1335 fn test_arrow_variant_to_literal_is_unsupported() {
1336 let variant_child = Arc::new(StructArray::from(vec![
1339 (
1340 Arc::new(Field::new("metadata", DataType::Binary, false)),
1341 Arc::new(BinaryArray::from(vec![Some(b"m".as_ref())])) as ArrayRef,
1342 ),
1343 (
1344 Arc::new(Field::new("value", DataType::Binary, false)),
1345 Arc::new(BinaryArray::from(vec![Some(b"v".as_ref())])) as ArrayRef,
1346 ),
1347 ])) as ArrayRef;
1348
1349 let struct_array = Arc::new(StructArray::from(vec![(
1350 Arc::new(
1351 Field::new("v", variant_child.data_type().clone(), false).with_metadata(
1352 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "1".to_string())]),
1353 ),
1354 ),
1355 variant_child,
1356 )])) as ArrayRef;
1357
1358 let ty = StructType::new(vec![
1359 NestedField::required(1, "v", Type::Variant(VariantType)).into(),
1360 ]);
1361
1362 let err = arrow_struct_to_literal(&struct_array, &ty).unwrap_err();
1363 assert_eq!(err.kind(), ErrorKind::FeatureUnsupported);
1364 assert!(
1365 err.to_string()
1366 .contains("Converting variant Arrow array to Iceberg literal is not supported yet"),
1367 "{err}"
1368 );
1369 }
1370
1371 #[test]
1372 fn test_find_field_by_id() {
1373 let field_a_array = Int32Array::from(vec![Some(42), Some(43), None]);
1375 let field_b_array = StringArray::from(vec![Some("value1"), Some("value2"), None]);
1376
1377 let nested_struct_array =
1379 Arc::new(StructArray::from(vec![
1380 (
1381 Arc::new(Field::new("field_a", DataType::Int32, true).with_metadata(
1382 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "1".to_string())]),
1383 )),
1384 Arc::new(field_a_array) as ArrayRef,
1385 ),
1386 (
1387 Arc::new(Field::new("field_b", DataType::Utf8, true).with_metadata(
1388 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "2".to_string())]),
1389 )),
1390 Arc::new(field_b_array) as ArrayRef,
1391 ),
1392 ])) as ArrayRef;
1393
1394 let field_c_array = Int32Array::from(vec![Some(100), Some(200), None]);
1395
1396 let struct_array = Arc::new(StructArray::from(vec![
1398 (
1399 Arc::new(
1400 Field::new(
1401 "nested_struct",
1402 DataType::Struct(Fields::from(vec![
1403 Field::new("field_a", DataType::Int32, true).with_metadata(
1404 HashMap::from([(
1405 PARQUET_FIELD_ID_META_KEY.to_string(),
1406 "1".to_string(),
1407 )]),
1408 ),
1409 Field::new("field_b", DataType::Utf8, true).with_metadata(
1410 HashMap::from([(
1411 PARQUET_FIELD_ID_META_KEY.to_string(),
1412 "2".to_string(),
1413 )]),
1414 ),
1415 ])),
1416 true,
1417 )
1418 .with_metadata(HashMap::from([(
1419 PARQUET_FIELD_ID_META_KEY.to_string(),
1420 "3".to_string(),
1421 )])),
1422 ),
1423 nested_struct_array,
1424 ),
1425 (
1426 Arc::new(Field::new("field_c", DataType::Int32, true).with_metadata(
1427 HashMap::from([(PARQUET_FIELD_ID_META_KEY.to_string(), "4".to_string())]),
1428 )),
1429 Arc::new(field_c_array) as ArrayRef,
1430 ),
1431 ])) as ArrayRef;
1432
1433 let accessor = ArrowArrayAccessor::new_with_match_mode(FieldMatchMode::Id);
1435
1436 let nested_field = NestedField::optional(
1438 3,
1439 "nested_struct",
1440 Type::Struct(StructType::new(vec![
1441 Arc::new(NestedField::optional(
1442 1,
1443 "field_a",
1444 Type::Primitive(PrimitiveType::Int),
1445 )),
1446 Arc::new(NestedField::optional(
1447 2,
1448 "field_b",
1449 Type::Primitive(PrimitiveType::String),
1450 )),
1451 ])),
1452 );
1453 let nested_partner = accessor
1454 .field_partner(&struct_array, &nested_field)
1455 .unwrap();
1456
1457 let field_a = NestedField::optional(1, "field_a", Type::Primitive(PrimitiveType::Int));
1459 let field_a_partner = accessor.field_partner(nested_partner, &field_a).unwrap();
1460
1461 let int_array = field_a_partner
1463 .as_any()
1464 .downcast_ref::<Int32Array>()
1465 .unwrap();
1466 assert_eq!(int_array.value(0), 42);
1467 assert_eq!(int_array.value(1), 43);
1468 assert!(int_array.is_null(2));
1469 }
1470
1471 #[test]
1472 fn test_find_field_by_name() {
1473 let field_a_array = Int32Array::from(vec![Some(42), Some(43), None]);
1475 let field_b_array = StringArray::from(vec![Some("value1"), Some("value2"), None]);
1476
1477 let nested_struct_array = Arc::new(StructArray::from(vec![
1479 (
1480 Arc::new(Field::new("field_a", DataType::Int32, true)),
1481 Arc::new(field_a_array) as ArrayRef,
1482 ),
1483 (
1484 Arc::new(Field::new("field_b", DataType::Utf8, true)),
1485 Arc::new(field_b_array) as ArrayRef,
1486 ),
1487 ])) as ArrayRef;
1488
1489 let field_c_array = Int32Array::from(vec![Some(100), Some(200), None]);
1490
1491 let struct_array = Arc::new(StructArray::from(vec![
1493 (
1494 Arc::new(Field::new(
1495 "nested_struct",
1496 DataType::Struct(Fields::from(vec![
1497 Field::new("field_a", DataType::Int32, true),
1498 Field::new("field_b", DataType::Utf8, true),
1499 ])),
1500 true,
1501 )),
1502 nested_struct_array,
1503 ),
1504 (
1505 Arc::new(Field::new("field_c", DataType::Int32, true)),
1506 Arc::new(field_c_array) as ArrayRef,
1507 ),
1508 ])) as ArrayRef;
1509
1510 let accessor = ArrowArrayAccessor::new_with_match_mode(FieldMatchMode::Name);
1512
1513 let nested_field = NestedField::optional(
1515 3,
1516 "nested_struct",
1517 Type::Struct(StructType::new(vec![
1518 Arc::new(NestedField::optional(
1519 1,
1520 "field_a",
1521 Type::Primitive(PrimitiveType::Int),
1522 )),
1523 Arc::new(NestedField::optional(
1524 2,
1525 "field_b",
1526 Type::Primitive(PrimitiveType::String),
1527 )),
1528 ])),
1529 );
1530 let nested_partner = accessor
1531 .field_partner(&struct_array, &nested_field)
1532 .unwrap();
1533
1534 let field_a = NestedField::optional(1, "field_a", Type::Primitive(PrimitiveType::Int));
1536 let field_a_partner = accessor.field_partner(nested_partner, &field_a).unwrap();
1537
1538 let int_array = field_a_partner
1540 .as_any()
1541 .downcast_ref::<Int32Array>()
1542 .unwrap();
1543 assert_eq!(int_array.value(0), 42);
1544 assert_eq!(int_array.value(1), 43);
1545 assert!(int_array.is_null(2));
1546 }
1547
1548 #[test]
1549 fn test_complex_nested() {
1550 let struct_type = StructType::new(vec![
1557 Arc::new(NestedField::required(
1558 0,
1559 "A",
1560 Type::List(ListType::new(Arc::new(NestedField::required(
1561 1,
1562 "item",
1563 Type::Struct(StructType::new(vec![
1564 Arc::new(NestedField::required(
1565 2,
1566 "a1",
1567 Type::Primitive(PrimitiveType::Int),
1568 )),
1569 Arc::new(NestedField::required(
1570 3,
1571 "a2",
1572 Type::Primitive(PrimitiveType::Int),
1573 )),
1574 ])),
1575 )))),
1576 )),
1577 Arc::new(NestedField::required(
1578 4,
1579 "B",
1580 Type::List(ListType::new(Arc::new(NestedField::required(
1581 5,
1582 "item",
1583 Type::Map(MapType::new(
1584 NestedField::optional(6, "keys", Type::Primitive(PrimitiveType::Int))
1585 .into(),
1586 NestedField::optional(7, "values", Type::Primitive(PrimitiveType::Int))
1587 .into(),
1588 )),
1589 )))),
1590 )),
1591 Arc::new(NestedField::required(
1592 8,
1593 "C",
1594 Type::List(ListType::new(Arc::new(NestedField::required(
1595 9,
1596 "item",
1597 Type::List(ListType::new(Arc::new(NestedField::optional(
1598 10,
1599 "item",
1600 Type::Primitive(PrimitiveType::Int),
1601 )))),
1602 )))),
1603 )),
1604 ]);
1605
1606 let struct_array =
1612 {
1613 let a_struct_a1_builder = Int32Builder::new();
1614 let a_struct_a2_builder = Int32Builder::new();
1615 let a_struct_builder =
1616 StructBuilder::new(
1617 vec![
1618 Field::new("a1", DataType::Int32, false).with_metadata(HashMap::from(
1619 [(PARQUET_FIELD_ID_META_KEY.to_string(), "2".to_string())],
1620 )),
1621 Field::new("a2", DataType::Int32, false).with_metadata(HashMap::from(
1622 [(PARQUET_FIELD_ID_META_KEY.to_string(), "3".to_string())],
1623 )),
1624 ],
1625 vec![Box::new(a_struct_a1_builder), Box::new(a_struct_a2_builder)],
1626 );
1627 let a_builder = ListBuilder::new(a_struct_builder);
1628
1629 let map_key_builder = Int32Builder::new();
1630 let map_value_builder = Int32Builder::new();
1631 let map_builder = MapBuilder::new(None, map_key_builder, map_value_builder);
1632 let b_builder = ListBuilder::new(map_builder);
1633
1634 let inner_list_item_builder = Int32Builder::new();
1635 let inner_list_builder = ListBuilder::new(inner_list_item_builder);
1636 let c_builder = ListBuilder::new(inner_list_builder);
1637
1638 let mut top_struct_builder = {
1639 let a_struct_type =
1640 DataType::Struct(Fields::from(vec![
1641 Field::new("a1", DataType::Int32, false).with_metadata(HashMap::from(
1642 [(PARQUET_FIELD_ID_META_KEY.to_string(), "2".to_string())],
1643 )),
1644 Field::new("a2", DataType::Int32, false).with_metadata(HashMap::from(
1645 [(PARQUET_FIELD_ID_META_KEY.to_string(), "3".to_string())],
1646 )),
1647 ]));
1648 let a_type =
1649 DataType::List(Arc::new(Field::new("item", a_struct_type.clone(), true)));
1650
1651 let b_map_entry_struct = Field::new(
1652 "entries",
1653 DataType::Struct(Fields::from(vec![
1654 Field::new("keys", DataType::Int32, false),
1655 Field::new("values", DataType::Int32, true),
1656 ])),
1657 false,
1658 );
1659 let b_map_type =
1660 DataType::Map(Arc::new(b_map_entry_struct), false);
1661 let b_type =
1662 DataType::List(Arc::new(Field::new("item", b_map_type.clone(), true)));
1663
1664 let c_inner_list_type =
1665 DataType::List(Arc::new(Field::new("item", DataType::Int32, true)));
1666 let c_type = DataType::List(Arc::new(Field::new(
1667 "item",
1668 c_inner_list_type.clone(),
1669 true,
1670 )));
1671 StructBuilder::new(
1672 Fields::from(vec![
1673 Field::new("A", a_type.clone(), false).with_metadata(HashMap::from([
1674 (PARQUET_FIELD_ID_META_KEY.to_string(), "0".to_string()),
1675 ])),
1676 Field::new("B", b_type.clone(), false).with_metadata(HashMap::from([
1677 (PARQUET_FIELD_ID_META_KEY.to_string(), "4".to_string()),
1678 ])),
1679 Field::new("C", c_type.clone(), false).with_metadata(HashMap::from([
1680 (PARQUET_FIELD_ID_META_KEY.to_string(), "8".to_string()),
1681 ])),
1682 ]),
1683 vec![
1684 Box::new(a_builder),
1685 Box::new(b_builder),
1686 Box::new(c_builder),
1687 ],
1688 )
1689 };
1690
1691 {
1694 let a_builder = top_struct_builder
1695 .field_builder::<ListBuilder<StructBuilder>>(0)
1696 .unwrap();
1697 let struct_builder = a_builder.values();
1698 struct_builder
1699 .field_builder::<Int32Builder>(0)
1700 .unwrap()
1701 .append_value(10);
1702 struct_builder
1703 .field_builder::<Int32Builder>(1)
1704 .unwrap()
1705 .append_value(20);
1706 struct_builder.append(true);
1707 let struct_builder = a_builder.values();
1708 struct_builder
1709 .field_builder::<Int32Builder>(0)
1710 .unwrap()
1711 .append_value(11);
1712 struct_builder
1713 .field_builder::<Int32Builder>(1)
1714 .unwrap()
1715 .append_value(21);
1716 struct_builder.append(true);
1717 a_builder.append(true);
1718 }
1719 {
1720 let b_builder = top_struct_builder
1721 .field_builder::<ListBuilder<MapBuilder<Int32Builder, Int32Builder>>>(1)
1722 .unwrap();
1723 let map_builder = b_builder.values();
1724 map_builder.keys().append_value(1);
1725 map_builder.values().append_value(100);
1726 map_builder.keys().append_value(3);
1727 map_builder.values().append_value(300);
1728 map_builder.append(true).unwrap();
1729
1730 map_builder.keys().append_value(2);
1731 map_builder.values().append_value(200);
1732 map_builder.append(true).unwrap();
1733
1734 b_builder.append(true);
1735 }
1736 {
1737 let c_builder = top_struct_builder
1738 .field_builder::<ListBuilder<ListBuilder<Int32Builder>>>(2)
1739 .unwrap();
1740 let inner_list_builder = c_builder.values();
1741 inner_list_builder.values().append_value(100);
1742 inner_list_builder.values().append_value(101);
1743 inner_list_builder.values().append_value(102);
1744 inner_list_builder.append(true);
1745 let inner_list_builder = c_builder.values();
1746 inner_list_builder.values().append_value(200);
1747 inner_list_builder.values().append_value(201);
1748 inner_list_builder.append(true);
1749 c_builder.append(true);
1750 }
1751 top_struct_builder.append(true);
1752
1753 {
1756 let a_builder = top_struct_builder
1757 .field_builder::<ListBuilder<StructBuilder>>(0)
1758 .unwrap();
1759 let struct_builder = a_builder.values();
1760 struct_builder
1761 .field_builder::<Int32Builder>(0)
1762 .unwrap()
1763 .append_value(12);
1764 struct_builder
1765 .field_builder::<Int32Builder>(1)
1766 .unwrap()
1767 .append_value(22);
1768 struct_builder.append(true);
1769 let struct_builder = a_builder.values();
1770 struct_builder
1771 .field_builder::<Int32Builder>(0)
1772 .unwrap()
1773 .append_value(13);
1774 struct_builder
1775 .field_builder::<Int32Builder>(1)
1776 .unwrap()
1777 .append_value(23);
1778 struct_builder.append(true);
1779 a_builder.append(true);
1780 }
1781 {
1782 let b_builder = top_struct_builder
1783 .field_builder::<ListBuilder<MapBuilder<Int32Builder, Int32Builder>>>(1)
1784 .unwrap();
1785 let map_builder = b_builder.values();
1786 map_builder.keys().append_value(3);
1787 map_builder.values().append_value(300);
1788 map_builder.append(true).unwrap();
1789
1790 b_builder.append(true);
1791 }
1792 {
1793 let c_builder = top_struct_builder
1794 .field_builder::<ListBuilder<ListBuilder<Int32Builder>>>(2)
1795 .unwrap();
1796 let inner_list_builder = c_builder.values();
1797 inner_list_builder.values().append_value(300);
1798 inner_list_builder.values().append_value(301);
1799 inner_list_builder.values().append_value(302);
1800 inner_list_builder.append(true);
1801 let inner_list_builder = c_builder.values();
1802 inner_list_builder.values().append_value(400);
1803 inner_list_builder.values().append_value(401);
1804 inner_list_builder.append(true);
1805 c_builder.append(true);
1806 }
1807 top_struct_builder.append(true);
1808
1809 Arc::new(top_struct_builder.finish()) as ArrayRef
1810 };
1811
1812 let result = arrow_struct_to_literal(&struct_array, &struct_type).unwrap();
1813 assert_eq!(result, vec![
1814 Some(Literal::Struct(Struct::from_iter(vec![
1815 Some(Literal::List(vec![
1816 Some(Literal::Struct(Struct::from_iter(vec![
1817 Some(Literal::int(10)),
1818 Some(Literal::int(20)),
1819 ]))),
1820 Some(Literal::Struct(Struct::from_iter(vec![
1821 Some(Literal::int(11)),
1822 Some(Literal::int(21)),
1823 ]))),
1824 ])),
1825 Some(Literal::List(vec![
1826 Some(Literal::Map(Map::from_iter(vec![
1827 (Literal::int(1), Some(Literal::int(100))),
1828 (Literal::int(3), Some(Literal::int(300))),
1829 ]))),
1830 Some(Literal::Map(Map::from_iter(vec![(
1831 Literal::int(2),
1832 Some(Literal::int(200))
1833 ),]))),
1834 ])),
1835 Some(Literal::List(vec![
1836 Some(Literal::List(vec![
1837 Some(Literal::int(100)),
1838 Some(Literal::int(101)),
1839 Some(Literal::int(102)),
1840 ])),
1841 Some(Literal::List(vec![
1842 Some(Literal::int(200)),
1843 Some(Literal::int(201)),
1844 ])),
1845 ])),
1846 ]))),
1847 Some(Literal::Struct(Struct::from_iter(vec![
1848 Some(Literal::List(vec![
1849 Some(Literal::Struct(Struct::from_iter(vec![
1850 Some(Literal::int(12)),
1851 Some(Literal::int(22)),
1852 ]))),
1853 Some(Literal::Struct(Struct::from_iter(vec![
1854 Some(Literal::int(13)),
1855 Some(Literal::int(23)),
1856 ]))),
1857 ])),
1858 Some(Literal::List(vec![Some(Literal::Map(Map::from_iter(
1859 vec![(Literal::int(3), Some(Literal::int(300))),]
1860 ))),])),
1861 Some(Literal::List(vec![
1862 Some(Literal::List(vec![
1863 Some(Literal::int(300)),
1864 Some(Literal::int(301)),
1865 Some(Literal::int(302)),
1866 ])),
1867 Some(Literal::List(vec![
1868 Some(Literal::int(400)),
1869 Some(Literal::int(401)),
1870 ])),
1871 ])),
1872 ]))),
1873 ]);
1874 }
1875
1876 #[test]
1877 fn test_create_decimal_array_respects_precision() {
1878 let target_precision = 18u8;
1881 let target_scale = 10i8;
1882 let target_type = DataType::Decimal128(target_precision, target_scale);
1883 let value = PrimitiveLiteral::Int128(10000000000);
1884
1885 let array = create_primitive_array_single_element(&target_type, &Some(value))
1886 .expect("Failed to create decimal array");
1887
1888 match array.data_type() {
1889 DataType::Decimal128(precision, scale) => {
1890 assert_eq!(*precision, target_precision);
1891 assert_eq!(*scale, target_scale);
1892 }
1893 other => panic!("Expected Decimal128, got {other:?}"),
1894 }
1895 }
1896
1897 #[test]
1898 fn test_create_decimal_array_repeated_respects_precision() {
1899 let target_precision = 18u8;
1901 let target_scale = 10i8;
1902 let target_type = DataType::Decimal128(target_precision, target_scale);
1903 let value = PrimitiveLiteral::Int128(10000000000);
1904 let num_rows = 5;
1905
1906 let array = create_primitive_array_repeated(&target_type, &Some(value), num_rows)
1907 .expect("Failed to create repeated decimal array");
1908
1909 match array.data_type() {
1910 DataType::Decimal128(precision, scale) => {
1911 assert_eq!(*precision, target_precision);
1912 assert_eq!(*scale, target_scale);
1913 }
1914 other => panic!("Expected Decimal128, got {other:?}"),
1915 }
1916
1917 assert_eq!(array.len(), num_rows);
1918 }
1919
1920 #[test]
1921 fn test_create_timestamp_microsecond_array_repeated() {
1922 let target_type = DataType::Timestamp(TimeUnit::Microsecond, None);
1923 let value = PrimitiveLiteral::Long(1_740_600_000_000_000);
1924 let num_rows = 3;
1925
1926 let array = create_primitive_array_repeated(&target_type, &Some(value), num_rows)
1927 .expect("Failed to create repeated timestamp microsecond array");
1928
1929 assert_eq!(array.data_type(), &target_type);
1930 assert_eq!(array.len(), num_rows);
1931 }
1932
1933 #[test]
1934 fn test_create_timestamp_microsecond_with_timezone_array_repeated() {
1935 let target_type = DataType::Timestamp(TimeUnit::Microsecond, Some("UTC".into()));
1936 let value = PrimitiveLiteral::Long(1_740_600_000_000_000);
1937 let num_rows = 2;
1938
1939 let array = create_primitive_array_repeated(&target_type, &Some(value), num_rows)
1940 .expect("Failed to create repeated timestamp microsecond array with timezone");
1941
1942 assert_eq!(array.data_type(), &target_type);
1943 assert_eq!(array.len(), num_rows);
1944 }
1945}