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