1use crate::codec::{AvroFieldBuilder, Tz};
482use crate::errors::AvroError;
483use crate::reader::header::read_header;
484use crate::schema::{
485 AvroSchema, CONFLUENT_MAGIC, Fingerprint, FingerprintAlgorithm, SCHEMA_METADATA_KEY,
486 SINGLE_OBJECT_MAGIC, Schema, SchemaStore,
487};
488use arrow_array::{RecordBatch, RecordBatchReader};
489use arrow_schema::{ArrowError, SchemaRef};
490use block::BlockDecoder;
491use header::Header;
492use indexmap::IndexMap;
493use record::RecordDecoder;
494use std::io::BufRead;
495
496mod block;
497mod cursor;
498mod header;
499mod record;
500mod vlq;
501
502#[cfg(feature = "async")]
503pub mod async_reader;
504
505pub use header::{HeaderInfo, read_header_info};
506
507#[expect(deprecated)]
508#[cfg(feature = "object_store")]
509pub use async_reader::AvroObjectReader;
510#[cfg(feature = "async")]
511pub use async_reader::{AsyncAvroFileReader, AsyncFileReader, SpawnedReader};
512
513fn is_incomplete_data(err: &AvroError) -> bool {
514 matches!(
515 err,
516 AvroError::EOF(_) | AvroError::NeedMoreData(_) | AvroError::NeedMoreDataRange(_)
517 )
518}
519
520#[derive(Debug)]
643pub struct Decoder {
644 active_decoder: RecordDecoder,
645 active_fingerprint: Option<Fingerprint>,
646 batch_size: usize,
647 remaining_capacity: usize,
648 cache: IndexMap<Fingerprint, RecordDecoder>,
649 fingerprint_algorithm: FingerprintAlgorithm,
650 pending_schema: Option<(Fingerprint, RecordDecoder)>,
651 awaiting_body: bool,
652}
653
654impl Decoder {
655 pub(crate) fn from_parts(
656 batch_size: usize,
657 active_decoder: RecordDecoder,
658 active_fingerprint: Option<Fingerprint>,
659 cache: IndexMap<Fingerprint, RecordDecoder>,
660 fingerprint_algorithm: FingerprintAlgorithm,
661 ) -> Self {
662 Self {
663 batch_size,
664 remaining_capacity: batch_size,
665 active_fingerprint,
666 active_decoder,
667 cache,
668 fingerprint_algorithm,
669 pending_schema: None,
670 awaiting_body: false,
671 }
672 }
673
674 pub fn schema(&self) -> SchemaRef {
679 self.active_decoder.schema().clone()
680 }
681
682 pub fn batch_size(&self) -> usize {
684 self.batch_size
685 }
686
687 pub fn decode(&mut self, data: &[u8]) -> Result<usize, AvroError> {
708 let mut total_consumed = 0usize;
709 while total_consumed < data.len() && self.remaining_capacity > 0 {
710 if self.awaiting_body {
711 match self.active_decoder.decode(&data[total_consumed..], 1) {
712 Ok(n) => {
713 self.remaining_capacity -= 1;
714 total_consumed += n;
715 self.awaiting_body = false;
716 continue;
717 }
718 Err(ref e) if is_incomplete_data(e) => break,
719 Err(e) => return Err(e),
720 };
721 }
722 match self.handle_prefix(&data[total_consumed..])? {
723 Some(0) => break, Some(n) => {
725 total_consumed += n;
726 self.apply_pending_schema_if_batch_empty();
727 self.awaiting_body = true;
728 }
729 None => {
730 return Err(AvroError::ParseError(
731 "Missing magic bytes and fingerprint".to_string(),
732 ));
733 }
734 }
735 }
736 Ok(total_consumed)
737 }
738
739 fn handle_prefix(&mut self, buf: &[u8]) -> Result<Option<usize>, AvroError> {
744 match self.fingerprint_algorithm {
745 FingerprintAlgorithm::Rabin => {
746 self.handle_prefix_common(buf, &SINGLE_OBJECT_MAGIC, |bytes| {
747 Fingerprint::Rabin(u64::from_le_bytes(bytes))
748 })
749 }
750 FingerprintAlgorithm::Id => self.handle_prefix_common(buf, &CONFLUENT_MAGIC, |bytes| {
751 Fingerprint::Id(u32::from_be_bytes(bytes))
752 }),
753 FingerprintAlgorithm::Id64 => {
754 self.handle_prefix_common(buf, &CONFLUENT_MAGIC, |bytes| {
755 Fingerprint::Id64(u64::from_be_bytes(bytes))
756 })
757 }
758 #[cfg(feature = "md5")]
759 FingerprintAlgorithm::MD5 => {
760 self.handle_prefix_common(buf, &SINGLE_OBJECT_MAGIC, |bytes| {
761 Fingerprint::MD5(bytes)
762 })
763 }
764 #[cfg(feature = "sha256")]
765 FingerprintAlgorithm::SHA256 => {
766 self.handle_prefix_common(buf, &SINGLE_OBJECT_MAGIC, |bytes| {
767 Fingerprint::SHA256(bytes)
768 })
769 }
770 }
771 }
772
773 fn handle_prefix_common<const MAGIC_LEN: usize, const N: usize>(
777 &mut self,
778 buf: &[u8],
779 magic: &[u8; MAGIC_LEN],
780 fingerprint_from: impl FnOnce([u8; N]) -> Fingerprint,
781 ) -> Result<Option<usize>, AvroError> {
782 if buf.len() < MAGIC_LEN {
785 return Ok(Some(0));
786 }
787 if &buf[..MAGIC_LEN] != magic {
789 return Ok(None);
790 }
791 let consumed_fp = self.handle_fingerprint(&buf[MAGIC_LEN..], fingerprint_from)?;
793 Ok(Some(consumed_fp.map_or(0, |n| n + MAGIC_LEN)))
796 }
797
798 fn handle_fingerprint<const N: usize>(
803 &mut self,
804 buf: &[u8],
805 fingerprint_from: impl FnOnce([u8; N]) -> Fingerprint,
806 ) -> Result<Option<usize>, AvroError> {
807 let Some(fingerprint_bytes) = buf.get(..N) else {
809 return Ok(None); };
811 let new_fingerprint = fingerprint_from(fingerprint_bytes.try_into().unwrap());
813 if self.active_fingerprint != Some(new_fingerprint) {
815 let Some(new_decoder) = self.cache.shift_remove(&new_fingerprint) else {
816 return Err(AvroError::ParseError(format!(
817 "Unknown fingerprint: {new_fingerprint:?}"
818 )));
819 };
820 self.pending_schema = Some((new_fingerprint, new_decoder));
821 if self.remaining_capacity < self.batch_size {
824 self.remaining_capacity = 0;
825 }
826 }
827 Ok(Some(N))
828 }
829
830 fn apply_pending_schema(&mut self) {
831 if let Some((new_fingerprint, new_decoder)) = self.pending_schema.take() {
832 if let Some(old_fingerprint) = self.active_fingerprint.replace(new_fingerprint) {
833 let old_decoder = std::mem::replace(&mut self.active_decoder, new_decoder);
834 self.cache.shift_remove(&old_fingerprint);
835 self.cache.insert(old_fingerprint, old_decoder);
836 } else {
837 self.active_decoder = new_decoder;
838 }
839 }
840 }
841
842 fn apply_pending_schema_if_batch_empty(&mut self) {
843 if self.batch_is_empty() {
844 self.apply_pending_schema();
845 }
846 }
847
848 fn flush_and_reset(&mut self) -> Result<Option<RecordBatch>, AvroError> {
849 if self.batch_is_empty() {
850 return Ok(None);
851 }
852 let batch = self.active_decoder.flush()?;
853 self.remaining_capacity = self.batch_size;
854 Ok(Some(batch))
855 }
856
857 pub fn flush(&mut self) -> Result<Option<RecordBatch>, AvroError> {
864 let batch = self.flush_and_reset();
866 self.apply_pending_schema();
867 batch
868 }
869
870 pub fn capacity(&self) -> usize {
872 self.remaining_capacity
873 }
874
875 pub fn batch_is_full(&self) -> bool {
877 self.remaining_capacity == 0
878 }
879
880 pub fn batch_is_empty(&self) -> bool {
882 self.remaining_capacity == self.batch_size
883 }
884
885 fn decode_block(&mut self, data: &[u8], count: usize) -> Result<(usize, usize), AvroError> {
889 let to_decode = std::cmp::min(count, self.remaining_capacity);
891 if to_decode == 0 {
892 return Ok((0, 0));
893 }
894 let consumed = self.active_decoder.decode(data, to_decode)?;
895 self.remaining_capacity -= to_decode;
896 Ok((consumed, to_decode))
897 }
898
899 fn flush_block(&mut self) -> Result<Option<RecordBatch>, AvroError> {
902 self.flush_and_reset()
903 }
904}
905
906#[derive(Debug)]
969pub struct ReaderBuilder {
970 batch_size: usize,
971 strict_mode: bool,
972 utf8_view: bool,
973 tz: Tz,
974 reader_schema: Option<AvroSchema>,
975 projection: Option<Vec<usize>>,
976 writer_schema_store: Option<SchemaStore>,
977 active_fingerprint: Option<Fingerprint>,
978}
979
980impl Default for ReaderBuilder {
981 fn default() -> Self {
982 Self {
983 batch_size: 1024,
984 strict_mode: false,
985 utf8_view: false,
986 tz: Default::default(),
987 reader_schema: None,
988 projection: None,
989 writer_schema_store: None,
990 active_fingerprint: None,
991 }
992 }
993}
994
995impl ReaderBuilder {
996 pub fn new() -> Self {
1007 Self::default()
1008 }
1009
1010 fn make_record_decoder(
1011 &self,
1012 writer_schema: &Schema,
1013 reader_schema: Option<&Schema>,
1014 ) -> Result<RecordDecoder, AvroError> {
1015 let mut builder = AvroFieldBuilder::new(writer_schema);
1016 if let Some(reader_schema) = reader_schema {
1017 builder = builder.with_reader_schema(reader_schema);
1018 }
1019 let root = builder
1020 .with_utf8view(self.utf8_view)
1021 .with_strict_mode(self.strict_mode)
1022 .with_tz(self.tz)
1023 .build()?;
1024 RecordDecoder::try_new_with_options(root.data_type())
1025 }
1026
1027 fn make_record_decoder_from_schemas(
1028 &self,
1029 writer_schema: &Schema,
1030 reader_schema: Option<&AvroSchema>,
1031 ) -> Result<RecordDecoder, AvroError> {
1032 let reader_schema_raw = reader_schema.map(|s| s.schema()).transpose()?;
1033 self.make_record_decoder(writer_schema, reader_schema_raw.as_ref())
1034 }
1035
1036 fn make_decoder(
1037 &self,
1038 header: Option<&Header>,
1039 reader_schema: Option<&AvroSchema>,
1040 ) -> Result<Decoder, AvroError> {
1041 if let Some(hdr) = header {
1042 let writer_schema = hdr.schema()?.ok_or_else(|| {
1043 AvroError::ParseError("No Avro schema present in file header".into())
1044 })?;
1045 let projected_reader_schema = self
1046 .projection
1047 .as_deref()
1048 .map(|projection| {
1049 let base_schema = if let Some(reader_schema) = reader_schema {
1050 reader_schema.clone()
1051 } else {
1052 let raw = hdr.get(SCHEMA_METADATA_KEY).ok_or_else(|| {
1053 AvroError::ParseError(
1054 "No Avro schema present in file header".to_string(),
1055 )
1056 })?;
1057 let json_string = std::str::from_utf8(raw)
1058 .map_err(|e| {
1059 AvroError::ParseError(format!(
1060 "Invalid UTF-8 in Avro schema header: {e}"
1061 ))
1062 })?
1063 .to_string();
1064 AvroSchema::new(json_string)
1065 };
1066 base_schema.project(projection)
1067 })
1068 .transpose()?;
1069 let effective_reader_schema = projected_reader_schema.as_ref().or(reader_schema);
1070 let record_decoder =
1071 self.make_record_decoder_from_schemas(&writer_schema, effective_reader_schema)?;
1072 return Ok(Decoder::from_parts(
1073 self.batch_size,
1074 record_decoder,
1075 None,
1076 IndexMap::new(),
1077 FingerprintAlgorithm::Rabin,
1078 ));
1079 }
1080 let store = self.writer_schema_store.as_ref().ok_or_else(|| {
1081 AvroError::ParseError("Writer schema store required for raw Avro".into())
1082 })?;
1083 let fingerprints = store.fingerprints();
1084 if fingerprints.is_empty() {
1085 return Err(AvroError::ParseError(
1086 "Writer schema store must contain at least one schema".into(),
1087 ));
1088 }
1089 let start_fingerprint = self
1090 .active_fingerprint
1091 .or_else(|| fingerprints.first().copied())
1092 .ok_or_else(|| {
1093 AvroError::ParseError("Could not determine initial schema fingerprint".into())
1094 })?;
1095 let projection = self.projection.as_deref();
1096 let projected_reader_schema = match (projection, reader_schema) {
1097 (Some(projection), Some(reader_schema)) => Some(reader_schema.project(projection)?),
1098 _ => None,
1099 };
1100 let mut cache = IndexMap::with_capacity(fingerprints.len().saturating_sub(1));
1101 let mut active_decoder: Option<RecordDecoder> = None;
1102 for fingerprint in store.fingerprints() {
1103 let Some(avro_schema) = store.lookup(&fingerprint) else {
1104 return Err(AvroError::General(format!(
1105 "Fingerprint {fingerprint:?} not found in schema store",
1106 )));
1107 };
1108 let writer_schema = avro_schema.schema()?;
1109 let record_decoder = match projection {
1110 None => self.make_record_decoder_from_schemas(&writer_schema, reader_schema)?,
1111 Some(projection) => {
1112 if let Some(ref pruned_reader_schema) = projected_reader_schema {
1113 self.make_record_decoder_from_schemas(
1114 &writer_schema,
1115 Some(pruned_reader_schema),
1116 )?
1117 } else {
1118 let derived_reader_schema = avro_schema.project(projection)?;
1119 self.make_record_decoder_from_schemas(
1120 &writer_schema,
1121 Some(&derived_reader_schema),
1122 )?
1123 }
1124 }
1125 };
1126 if fingerprint == start_fingerprint {
1127 active_decoder = Some(record_decoder);
1128 } else {
1129 cache.insert(fingerprint, record_decoder);
1130 }
1131 }
1132 let active_decoder = active_decoder.ok_or_else(|| {
1133 AvroError::General(format!(
1134 "Initial fingerprint {start_fingerprint:?} not found in schema store"
1135 ))
1136 })?;
1137 Ok(Decoder::from_parts(
1138 self.batch_size,
1139 active_decoder,
1140 Some(start_fingerprint),
1141 cache,
1142 store.fingerprint_algorithm(),
1143 ))
1144 }
1145
1146 pub fn with_batch_size(mut self, batch_size: usize) -> Self {
1152 self.batch_size = batch_size;
1153 self
1154 }
1155
1156 pub fn with_utf8_view(mut self, utf8_view: bool) -> Self {
1162 self.utf8_view = utf8_view;
1163 self
1164 }
1165
1166 pub fn use_utf8view(&self) -> bool {
1168 self.utf8_view
1169 }
1170
1171 pub fn with_strict_mode(mut self, strict_mode: bool) -> Self {
1176 self.strict_mode = strict_mode;
1177 self
1178 }
1179
1180 pub fn with_tz(mut self, tz: Tz) -> Self {
1184 self.tz = tz;
1185 self
1186 }
1187
1188 pub fn with_reader_schema(mut self, schema: AvroSchema) -> Self {
1195 self.reader_schema = Some(schema);
1196 self
1197 }
1198
1199 pub fn with_projection(mut self, projection: Vec<usize>) -> Self {
1257 self.projection = Some(projection);
1258 self
1259 }
1260
1261 pub fn with_writer_schema_store(mut self, store: SchemaStore) -> Self {
1269 self.writer_schema_store = Some(store);
1270 self
1271 }
1272
1273 pub fn with_active_fingerprint(mut self, fp: Fingerprint) -> Self {
1278 self.active_fingerprint = Some(fp);
1279 self
1280 }
1281
1282 pub fn build<R: BufRead>(self, mut reader: R) -> Result<Reader<R>, ArrowError> {
1288 let (header, _) = read_header(&mut reader)?;
1289 let decoder = self.make_decoder(Some(&header), self.reader_schema.as_ref())?;
1290 Ok(Reader {
1291 reader,
1292 header,
1293 decoder,
1294 block_decoder: BlockDecoder::default(),
1295 block_data: Vec::new(),
1296 block_count: 0,
1297 block_cursor: 0,
1298 finished: false,
1299 })
1300 }
1301
1302 pub fn build_decoder(self) -> Result<Decoder, ArrowError> {
1311 if self.writer_schema_store.is_none() {
1312 return Err(ArrowError::InvalidArgumentError(
1313 "Building a decoder requires a writer schema store".to_string(),
1314 ));
1315 }
1316 self.make_decoder(None, self.reader_schema.as_ref())
1317 .map_err(ArrowError::from)
1318 }
1319}
1320
1321#[derive(Debug)]
1331pub struct Reader<R: BufRead> {
1332 reader: R,
1333 header: Header,
1334 decoder: Decoder,
1335 block_decoder: BlockDecoder,
1336 block_data: Vec<u8>,
1337 block_count: usize,
1338 block_cursor: usize,
1339 finished: bool,
1340}
1341
1342impl<R: BufRead> Reader<R> {
1343 pub fn schema(&self) -> SchemaRef {
1346 self.decoder.schema()
1347 }
1348
1349 pub fn avro_header(&self) -> &Header {
1351 &self.header
1352 }
1353
1354 fn read(&mut self) -> Result<Option<RecordBatch>, AvroError> {
1359 'outer: while !self.finished && !self.decoder.batch_is_full() {
1360 while self.block_cursor == self.block_data.len() {
1361 let buf = self.reader.fill_buf()?;
1362 if buf.is_empty() {
1363 self.finished = true;
1364 break 'outer;
1365 }
1366 let consumed = self.block_decoder.decode(buf)?;
1368 self.reader.consume(consumed);
1369 if let Some(block) = self.block_decoder.flush() {
1370 if block.sync != self.header.sync() {
1372 return Err(AvroError::ParseError(
1373 "Avro block sync marker does not match file header".to_string(),
1374 ));
1375 }
1376 self.block_data = if let Some(ref codec) = self.header.compression()? {
1377 let decompressed: Vec<u8> = codec.decompress(&block.data)?;
1378 decompressed
1379 } else {
1380 block.data
1381 };
1382 self.block_count = block.count;
1383 self.block_cursor = 0;
1384 } else if consumed == 0 {
1385 return Err(AvroError::ParseError(
1387 "Could not decode next Avro block from partial data".to_string(),
1388 ));
1389 }
1390 }
1391 if self.block_cursor < self.block_data.len() {
1393 let (consumed, records_decoded) = self
1394 .decoder
1395 .decode_block(&self.block_data[self.block_cursor..], self.block_count)?;
1396 self.block_cursor += consumed;
1397 self.block_count -= records_decoded;
1398 }
1399 }
1400 self.decoder.flush_block()
1401 }
1402}
1403
1404impl<R: BufRead> Iterator for Reader<R> {
1405 type Item = Result<RecordBatch, ArrowError>;
1406
1407 fn next(&mut self) -> Option<Self::Item> {
1408 self.read().map_err(ArrowError::from).transpose()
1409 }
1410}
1411
1412impl<R: BufRead> RecordBatchReader for Reader<R> {
1413 fn schema(&self) -> SchemaRef {
1414 self.schema()
1415 }
1416}
1417
1418#[cfg(test)]
1419mod test {
1420 use crate::codec::{AvroFieldBuilder, Tz};
1421 use crate::reader::header::HeaderDecoder;
1422 use crate::reader::record::RecordDecoder;
1423 use crate::reader::{Decoder, Reader, ReaderBuilder};
1424 use crate::schema::{
1425 AVRO_ENUM_SYMBOLS_METADATA_KEY, AVRO_NAME_METADATA_KEY, AVRO_NAMESPACE_METADATA_KEY,
1426 AvroSchema, CONFLUENT_MAGIC, Fingerprint, FingerprintAlgorithm, PrimitiveType,
1427 SINGLE_OBJECT_MAGIC, SchemaStore,
1428 };
1429 use crate::test_util::arrow_test_data;
1430 use crate::writer::AvroWriter;
1431 use arrow_array::builder::{
1432 ArrayBuilder, BooleanBuilder, Float32Builder, Int32Builder, Int64Builder, ListBuilder,
1433 MapBuilder, StringBuilder, StructBuilder,
1434 };
1435 #[cfg(feature = "snappy")]
1436 use arrow_array::builder::{Float64Builder, MapFieldNames};
1437 use arrow_array::cast::AsArray;
1438 #[cfg(not(feature = "avro_custom_types"))]
1439 use arrow_array::types::Int64Type;
1440 #[cfg(feature = "avro_custom_types")]
1441 use arrow_array::types::{
1442 DurationMicrosecondType, DurationMillisecondType, DurationNanosecondType,
1443 DurationSecondType,
1444 };
1445 use arrow_array::types::{Int32Type, IntervalMonthDayNanoType};
1446 use arrow_array::*;
1447 #[cfg(feature = "snappy")]
1448 use arrow_buffer::{Buffer, NullBuffer};
1449 use arrow_buffer::{IntervalMonthDayNano, OffsetBuffer, ScalarBuffer, i256};
1450 #[cfg(feature = "avro_custom_types")]
1451 use arrow_schema::{
1452 ArrowError, DataType, Field, FieldRef, Fields, IntervalUnit, Schema, TimeUnit, UnionFields,
1453 UnionMode,
1454 };
1455 #[cfg(not(feature = "avro_custom_types"))]
1456 use arrow_schema::{
1457 ArrowError, DataType, Field, FieldRef, Fields, IntervalUnit, Schema, UnionFields, UnionMode,
1458 };
1459 use bytes::Bytes;
1460 use futures::executor::block_on;
1461 use futures::{Stream, StreamExt, TryStreamExt, stream};
1462 use serde_json::{Value, json};
1463 use std::collections::HashMap;
1464 use std::fs::File;
1465 use std::io::{BufReader, Cursor};
1466 use std::sync::Arc;
1467
1468 fn files() -> impl Iterator<Item = &'static str> {
1469 [
1470 #[cfg(feature = "snappy")]
1472 "avro/alltypes_plain.avro",
1473 #[cfg(all(feature = "snappy", not(miri)))]
1475 "avro/alltypes_plain.snappy.avro",
1476 #[cfg(all(feature = "zstd", not(miri)))]
1477 "avro/alltypes_plain.zstandard.avro",
1478 #[cfg(all(feature = "bzip2", not(miri)))]
1479 "avro/alltypes_plain.bzip2.avro",
1480 #[cfg(all(feature = "xz", not(miri)))]
1481 "avro/alltypes_plain.xz.avro",
1482 ]
1483 .into_iter()
1484 }
1485
1486 fn read_file(path: &str, batch_size: usize, utf8_view: bool) -> RecordBatch {
1487 let file = File::open(path).unwrap();
1488 let reader = ReaderBuilder::new()
1489 .with_batch_size(batch_size)
1490 .with_utf8_view(utf8_view)
1491 .build(BufReader::new(file))
1492 .unwrap();
1493 let schema = reader.schema();
1494 let batches = reader.collect::<Result<Vec<_>, _>>().unwrap();
1495 arrow::compute::concat_batches(&schema, &batches).unwrap()
1496 }
1497
1498 #[test]
1499 fn test_block_sync_marker_mismatch_errors() {
1500 let path = arrow_test_data("avro/alltypes_plain.avro");
1501 let mut bytes = std::fs::read(&path).unwrap();
1502 let last = bytes.len() - 1;
1504 bytes[last] ^= 0xFF;
1505 let reader = ReaderBuilder::new()
1506 .with_batch_size(1024)
1507 .build(std::io::Cursor::new(bytes))
1508 .unwrap();
1509 let err = reader
1510 .collect::<Result<Vec<_>, _>>()
1511 .expect_err("corrupted block sync marker should fail the read");
1512 assert!(err.to_string().contains("sync marker"), "{err}");
1513 }
1514
1515 fn read_file_strict(
1516 path: &str,
1517 batch_size: usize,
1518 utf8_view: bool,
1519 ) -> Result<Reader<BufReader<File>>, ArrowError> {
1520 let file = File::open(path)?;
1521 ReaderBuilder::new()
1522 .with_batch_size(batch_size)
1523 .with_utf8_view(utf8_view)
1524 .with_strict_mode(true)
1525 .build(BufReader::new(file))
1526 }
1527
1528 fn decode_stream<S: Stream<Item = Bytes> + Unpin>(
1529 mut decoder: Decoder,
1530 mut input: S,
1531 ) -> impl Stream<Item = Result<RecordBatch, ArrowError>> {
1532 async_stream::try_stream! {
1533 if let Some(data) = input.next().await {
1534 let consumed = decoder.decode(&data)?;
1535 if consumed < data.len() {
1536 Err(ArrowError::ParseError(
1537 "did not consume all bytes".to_string(),
1538 ))?;
1539 }
1540 }
1541 if let Some(batch) = decoder.flush()? {
1542 yield batch
1543 }
1544 }
1545 }
1546
1547 fn make_record_schema(pt: PrimitiveType) -> AvroSchema {
1548 let js = format!(
1549 r#"{{"type":"record","name":"TestRecord","fields":[{{"name":"a","type":"{}"}}]}}"#,
1550 pt.as_ref()
1551 );
1552 AvroSchema::new(js)
1553 }
1554
1555 fn make_two_schema_store() -> (
1556 SchemaStore,
1557 Fingerprint,
1558 Fingerprint,
1559 AvroSchema,
1560 AvroSchema,
1561 ) {
1562 let schema_int = make_record_schema(PrimitiveType::Int);
1563 let schema_long = make_record_schema(PrimitiveType::Long);
1564 let mut store = SchemaStore::new();
1565 let fp_int = store
1566 .register(schema_int.clone())
1567 .expect("register int schema");
1568 let fp_long = store
1569 .register(schema_long.clone())
1570 .expect("register long schema");
1571 (store, fp_int, fp_long, schema_int, schema_long)
1572 }
1573
1574 fn make_prefix(fp: Fingerprint) -> Vec<u8> {
1575 match fp {
1576 Fingerprint::Rabin(v) => {
1577 let mut out = Vec::with_capacity(2 + 8);
1578 out.extend_from_slice(&SINGLE_OBJECT_MAGIC);
1579 out.extend_from_slice(&v.to_le_bytes());
1580 out
1581 }
1582 Fingerprint::Id(v) => {
1583 panic!("make_prefix expects a Rabin fingerprint, got ({v})");
1584 }
1585 Fingerprint::Id64(v) => {
1586 panic!("make_prefix expects a Rabin fingerprint, got ({v})");
1587 }
1588 #[cfg(feature = "md5")]
1589 Fingerprint::MD5(v) => {
1590 panic!("make_prefix expects a Rabin fingerprint, got ({v:?})");
1591 }
1592 #[cfg(feature = "sha256")]
1593 Fingerprint::SHA256(id) => {
1594 panic!("make_prefix expects a Rabin fingerprint, got ({id:?})");
1595 }
1596 }
1597 }
1598
1599 fn make_decoder(store: &SchemaStore, fp: Fingerprint, reader_schema: &AvroSchema) -> Decoder {
1600 ReaderBuilder::new()
1601 .with_batch_size(8)
1602 .with_reader_schema(reader_schema.clone())
1603 .with_writer_schema_store(store.clone())
1604 .with_active_fingerprint(fp)
1605 .build_decoder()
1606 .expect("decoder")
1607 }
1608
1609 fn make_id_prefix(id: u32, additional: usize) -> Vec<u8> {
1610 let capacity = CONFLUENT_MAGIC.len() + size_of::<u32>() + additional;
1611 let mut out = Vec::with_capacity(capacity);
1612 out.extend_from_slice(&CONFLUENT_MAGIC);
1613 out.extend_from_slice(&id.to_be_bytes());
1614 out
1615 }
1616
1617 fn make_message_id(id: u32, value: i64) -> Vec<u8> {
1618 let encoded_value = encode_zigzag(value);
1619 let mut msg = make_id_prefix(id, encoded_value.len());
1620 msg.extend_from_slice(&encoded_value);
1621 msg
1622 }
1623
1624 fn make_id64_prefix(id: u64, additional: usize) -> Vec<u8> {
1625 let capacity = CONFLUENT_MAGIC.len() + size_of::<u64>() + additional;
1626 let mut out = Vec::with_capacity(capacity);
1627 out.extend_from_slice(&CONFLUENT_MAGIC);
1628 out.extend_from_slice(&id.to_be_bytes());
1629 out
1630 }
1631
1632 fn make_message_id64(id: u64, value: i64) -> Vec<u8> {
1633 let encoded_value = encode_zigzag(value);
1634 let mut msg = make_id64_prefix(id, encoded_value.len());
1635 msg.extend_from_slice(&encoded_value);
1636 msg
1637 }
1638
1639 fn make_value_schema(pt: PrimitiveType) -> AvroSchema {
1640 let json_schema = format!(
1641 r#"{{"type":"record","name":"S","fields":[{{"name":"v","type":"{}"}}]}}"#,
1642 pt.as_ref()
1643 );
1644 AvroSchema::new(json_schema)
1645 }
1646
1647 fn encode_zigzag(value: i64) -> Vec<u8> {
1648 let mut n = ((value << 1) ^ (value >> 63)) as u64;
1649 let mut out = Vec::new();
1650 loop {
1651 if (n & !0x7F) == 0 {
1652 out.push(n as u8);
1653 break;
1654 } else {
1655 out.push(((n & 0x7F) | 0x80) as u8);
1656 n >>= 7;
1657 }
1658 }
1659 out
1660 }
1661
1662 fn make_message(fp: Fingerprint, value: i64) -> Vec<u8> {
1663 let mut msg = make_prefix(fp);
1664 msg.extend_from_slice(&encode_zigzag(value));
1665 msg
1666 }
1667
1668 fn load_writer_schema_json(path: &str) -> Value {
1669 let file = File::open(path).unwrap();
1670 let (header, _) = super::read_header(BufReader::new(file)).unwrap();
1671 let schema = header.schema().unwrap().unwrap();
1672 serde_json::to_value(&schema).unwrap()
1673 }
1674
1675 fn make_reader_schema_with_promotions(
1676 path: &str,
1677 promotions: &HashMap<&str, &str>,
1678 ) -> AvroSchema {
1679 let mut root = load_writer_schema_json(path);
1680 assert_eq!(root["type"], "record", "writer schema must be a record");
1681 let fields = root
1682 .get_mut("fields")
1683 .and_then(|f| f.as_array_mut())
1684 .expect("record has fields");
1685 for f in fields.iter_mut() {
1686 let Some(name) = f.get("name").and_then(|n| n.as_str()) else {
1687 continue;
1688 };
1689 if let Some(new_ty) = promotions.get(name) {
1690 let ty = f.get_mut("type").expect("field has a type");
1691 match ty {
1692 Value::String(_) => {
1693 *ty = Value::String((*new_ty).to_string());
1694 }
1695 Value::Array(arr) => {
1697 for b in arr.iter_mut() {
1698 match b {
1699 Value::String(s) if s != "null" => {
1700 *b = Value::String((*new_ty).to_string());
1701 break;
1702 }
1703 Value::Object(_) => {
1704 *b = Value::String((*new_ty).to_string());
1705 break;
1706 }
1707 _ => {}
1708 }
1709 }
1710 }
1711 Value::Object(_) => {
1712 *ty = Value::String((*new_ty).to_string());
1713 }
1714 _ => {}
1715 }
1716 }
1717 }
1718 AvroSchema::new(root.to_string())
1719 }
1720
1721 fn make_reader_schema_with_enum_remap(
1722 path: &str,
1723 remap: &HashMap<&str, Vec<&str>>,
1724 ) -> AvroSchema {
1725 let mut root = load_writer_schema_json(path);
1726 assert_eq!(root["type"], "record", "writer schema must be a record");
1727 let fields = root
1728 .get_mut("fields")
1729 .and_then(|f| f.as_array_mut())
1730 .expect("record has fields");
1731
1732 fn to_symbols_array(symbols: &[&str]) -> Value {
1733 Value::Array(symbols.iter().map(|s| Value::String((*s).into())).collect())
1734 }
1735
1736 fn update_enum_symbols(ty: &mut Value, symbols: &Value) {
1737 match ty {
1738 Value::Object(map) => {
1739 if matches!(map.get("type"), Some(Value::String(t)) if t == "enum") {
1740 map.insert("symbols".to_string(), symbols.clone());
1741 }
1742 }
1743 Value::Array(arr) => {
1744 for b in arr.iter_mut() {
1745 if let Value::Object(map) = b
1746 && matches!(map.get("type"), Some(Value::String(t)) if t == "enum")
1747 {
1748 map.insert("symbols".to_string(), symbols.clone());
1749 }
1750 }
1751 }
1752 _ => {}
1753 }
1754 }
1755 for f in fields.iter_mut() {
1756 let Some(name) = f.get("name").and_then(|n| n.as_str()) else {
1757 continue;
1758 };
1759 if let Some(new_symbols) = remap.get(name) {
1760 let symbols_val = to_symbols_array(new_symbols);
1761 let ty = f.get_mut("type").expect("field has a type");
1762 update_enum_symbols(ty, &symbols_val);
1763 }
1764 }
1765 AvroSchema::new(root.to_string())
1766 }
1767
1768 fn read_alltypes_with_reader_schema(path: &str, reader_schema: AvroSchema) -> RecordBatch {
1769 let file = File::open(path).unwrap();
1770 let reader = ReaderBuilder::new()
1771 .with_batch_size(1024)
1772 .with_utf8_view(false)
1773 .with_reader_schema(reader_schema)
1774 .build(BufReader::new(file))
1775 .unwrap();
1776 let schema = reader.schema();
1777 let batches = reader.collect::<Result<Vec<_>, _>>().unwrap();
1778 arrow::compute::concat_batches(&schema, &batches).unwrap()
1779 }
1780
1781 fn make_reader_schema_with_selected_fields_in_order(
1782 path: &str,
1783 selected: &[&str],
1784 ) -> AvroSchema {
1785 let mut root = load_writer_schema_json(path);
1786 assert_eq!(root["type"], "record", "writer schema must be a record");
1787 let writer_fields = root
1788 .get("fields")
1789 .and_then(|f| f.as_array())
1790 .expect("record has fields");
1791 let mut field_map: HashMap<String, Value> = HashMap::with_capacity(writer_fields.len());
1792 for f in writer_fields {
1793 if let Some(name) = f.get("name").and_then(|n| n.as_str()) {
1794 field_map.insert(name.to_string(), f.clone());
1795 }
1796 }
1797 let mut new_fields = Vec::with_capacity(selected.len());
1798 for name in selected {
1799 let f = field_map
1800 .get(*name)
1801 .unwrap_or_else(|| panic!("field '{name}' not found in writer schema"))
1802 .clone();
1803 new_fields.push(f);
1804 }
1805 root["fields"] = Value::Array(new_fields);
1806 AvroSchema::new(root.to_string())
1807 }
1808
1809 fn write_ocf(schema: &Schema, batches: &[RecordBatch]) -> Vec<u8> {
1810 let mut w = AvroWriter::new(Vec::<u8>::new(), schema.clone()).expect("writer");
1811 for b in batches {
1812 w.write(b).expect("write");
1813 }
1814 w.finish().expect("finish");
1815 w.into_inner()
1816 }
1817
1818 #[test]
1819 fn ocf_projection_no_reader_schema_reorder() -> Result<(), Box<dyn std::error::Error>> {
1820 let writer_schema = Schema::new(vec![
1822 Field::new("id", DataType::Int32, false),
1823 Field::new("name", DataType::Utf8, false),
1824 Field::new("is_active", DataType::Boolean, false),
1825 ]);
1826 let batch = RecordBatch::try_new(
1827 Arc::new(writer_schema.clone()),
1828 vec![
1829 Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef,
1830 Arc::new(StringArray::from(vec!["a", "b"])) as ArrayRef,
1831 Arc::new(BooleanArray::from(vec![true, false])) as ArrayRef,
1832 ],
1833 )?;
1834 let bytes = write_ocf(&writer_schema, &[batch]);
1835 let mut reader = ReaderBuilder::new()
1837 .with_projection(vec![2, 0])
1838 .build(Cursor::new(bytes))?;
1839 let out = reader.next().unwrap()?;
1840 assert_eq!(out.num_columns(), 2);
1841 assert_eq!(out.schema().field(0).name(), "is_active");
1842 assert_eq!(out.schema().field(1).name(), "id");
1843 let is_active = out.column(0).as_boolean();
1844 assert!(is_active.value(0));
1845 assert!(!is_active.value(1));
1846 let id = out.column(1).as_primitive::<Int32Type>();
1847 assert_eq!(id.value(0), 1);
1848 assert_eq!(id.value(1), 2);
1849 Ok(())
1850 }
1851
1852 #[test]
1853 fn ocf_projection_with_reader_schema_alias_and_default()
1854 -> Result<(), Box<dyn std::error::Error>> {
1855 let writer_schema = Schema::new(vec![
1857 Field::new("id", DataType::Int64, false),
1858 Field::new("name", DataType::Utf8, false),
1859 ]);
1860 let batch = RecordBatch::try_new(
1861 Arc::new(writer_schema.clone()),
1862 vec![
1863 Arc::new(Int64Array::from(vec![1, 2])) as ArrayRef,
1864 Arc::new(StringArray::from(vec!["a", "b"])) as ArrayRef,
1865 ],
1866 )?;
1867 let bytes = write_ocf(&writer_schema, &[batch]);
1868 let reader_json = r#"
1872 {
1873 "type": "record",
1874 "name": "topLevelRecord",
1875 "fields": [
1876 { "name": "id", "type": "long" },
1877 { "name": "full_name", "type": ["null","string"], "aliases": ["name"], "default": null },
1878 { "name": "is_active", "type": "boolean", "default": true }
1879 ]
1880 }"#;
1881 let mut reader = ReaderBuilder::new()
1883 .with_reader_schema(AvroSchema::new(reader_json.to_string()))
1884 .with_projection(vec![1, 2])
1885 .build(Cursor::new(bytes))?;
1886 let out = reader.next().unwrap()?;
1887 assert_eq!(out.num_columns(), 2);
1888 assert_eq!(out.schema().field(0).name(), "full_name");
1889 assert_eq!(out.schema().field(1).name(), "is_active");
1890 let full_name = out.column(0).as_string::<i32>();
1891 assert_eq!(full_name.value(0), "a");
1892 assert_eq!(full_name.value(1), "b");
1893 let is_active = out.column(1).as_boolean();
1894 assert!(is_active.value(0));
1895 assert!(is_active.value(1));
1896 Ok(())
1897 }
1898
1899 #[test]
1900 fn projection_errors_out_of_bounds_and_duplicate() -> Result<(), Box<dyn std::error::Error>> {
1901 let writer_schema = Schema::new(vec![
1902 Field::new("a", DataType::Int32, false),
1903 Field::new("b", DataType::Int32, false),
1904 ]);
1905 let batch = RecordBatch::try_new(
1906 Arc::new(writer_schema.clone()),
1907 vec![
1908 Arc::new(Int32Array::from(vec![1])) as ArrayRef,
1909 Arc::new(Int32Array::from(vec![2])) as ArrayRef,
1910 ],
1911 )?;
1912 let bytes = write_ocf(&writer_schema, &[batch]);
1913 let err = ReaderBuilder::new()
1914 .with_projection(vec![2])
1915 .build(Cursor::new(bytes.clone()))
1916 .unwrap_err();
1917 assert!(matches!(err, ArrowError::AvroError(_)));
1918 assert!(err.to_string().contains("out of bounds"));
1919 let err = ReaderBuilder::new()
1920 .with_projection(vec![0, 0])
1921 .build(Cursor::new(bytes))
1922 .unwrap_err();
1923 assert!(matches!(err, ArrowError::AvroError(_)));
1924 assert!(err.to_string().contains("Duplicate projection index"));
1925 Ok(())
1926 }
1927
1928 #[test]
1929 #[cfg(feature = "snappy")]
1930 fn test_alltypes_plain_with_projection_and_reader_schema() {
1931 use std::fs::File;
1932 use std::io::BufReader;
1933 let path = arrow_test_data("avro/alltypes_plain.avro");
1934 let reader_schema = make_reader_schema_with_selected_fields_in_order(
1936 &path,
1937 &["double_col", "id", "tinyint_col"],
1938 );
1939 let file = File::open(&path).expect("open avro/alltypes_plain.avro");
1940 let reader = ReaderBuilder::new()
1941 .with_batch_size(1024)
1942 .with_reader_schema(reader_schema)
1943 .with_projection(vec![1, 2]) .build(BufReader::new(file))
1945 .expect("build reader with projection and reader schema");
1946 let schema = reader.schema();
1947 assert_eq!(schema.fields().len(), 2);
1949 assert_eq!(schema.field(0).name(), "id");
1950 assert_eq!(schema.field(1).name(), "tinyint_col");
1951 let batches: Vec<RecordBatch> = reader.collect::<Result<Vec<_>, _>>().unwrap();
1952 assert_eq!(batches.len(), 1);
1953 let batch = &batches[0];
1954 assert_eq!(batch.num_rows(), 8);
1955 assert_eq!(batch.num_columns(), 2);
1956 let expected = RecordBatch::try_from_iter_with_nullable([
1960 (
1961 "id",
1962 Arc::new(Int32Array::from(vec![4, 5, 6, 7, 2, 3, 0, 1])) as ArrayRef,
1963 true,
1964 ),
1965 (
1966 "tinyint_col",
1967 Arc::new(Int32Array::from(vec![0, 1, 0, 1, 0, 1, 0, 1])) as ArrayRef,
1968 true,
1969 ),
1970 ])
1971 .unwrap();
1972 assert_eq!(
1973 batch, &expected,
1974 "Projected batch mismatch for alltypes_plain.avro with reader schema and projection [1, 2]"
1975 );
1976 }
1977
1978 #[test]
1979 #[cfg(feature = "snappy")]
1980 fn test_alltypes_plain_with_projection() {
1981 use std::fs::File;
1982 use std::io::BufReader;
1983 let path = arrow_test_data("avro/alltypes_plain.avro");
1984 let file = File::open(&path).expect("open avro/alltypes_plain.avro");
1985 let reader = ReaderBuilder::new()
1986 .with_batch_size(1024)
1987 .with_projection(vec![2, 0, 5])
1988 .build(BufReader::new(file))
1989 .expect("build reader with projection");
1990 let schema = reader.schema();
1991 assert_eq!(schema.fields().len(), 3);
1992 assert_eq!(schema.field(0).name(), "tinyint_col");
1993 assert_eq!(schema.field(1).name(), "id");
1994 assert_eq!(schema.field(2).name(), "bigint_col");
1995 let batches: Vec<RecordBatch> = reader.collect::<Result<Vec<_>, _>>().unwrap();
1996 assert_eq!(batches.len(), 1);
1997 let batch = &batches[0];
1998 assert_eq!(batch.num_rows(), 8);
1999 assert_eq!(batch.num_columns(), 3);
2000 let expected = RecordBatch::try_from_iter_with_nullable([
2001 (
2002 "tinyint_col",
2003 Arc::new(Int32Array::from(vec![0, 1, 0, 1, 0, 1, 0, 1])) as ArrayRef,
2004 true,
2005 ),
2006 (
2007 "id",
2008 Arc::new(Int32Array::from(vec![4, 5, 6, 7, 2, 3, 0, 1])) as ArrayRef,
2009 true,
2010 ),
2011 (
2012 "bigint_col",
2013 Arc::new(Int64Array::from(vec![0, 10, 0, 10, 0, 10, 0, 10])) as ArrayRef,
2014 true,
2015 ),
2016 ])
2017 .unwrap();
2018 assert_eq!(
2019 batch, &expected,
2020 "Projected batch mismatch for alltypes_plain.avro with projection [2, 0, 5]"
2021 );
2022 }
2023
2024 #[test]
2025 fn writer_string_reader_nullable_with_alias() -> Result<(), Box<dyn std::error::Error>> {
2026 let writer_schema = Schema::new(vec![
2027 Field::new("id", DataType::Int64, false),
2028 Field::new("name", DataType::Utf8, false),
2029 ]);
2030 let batch = RecordBatch::try_new(
2031 Arc::new(writer_schema.clone()),
2032 vec![
2033 Arc::new(Int64Array::from(vec![1, 2])) as ArrayRef,
2034 Arc::new(StringArray::from(vec!["a", "b"])) as ArrayRef,
2035 ],
2036 )?;
2037 let bytes = write_ocf(&writer_schema, &[batch]);
2038 let reader_json = r#"
2039 {
2040 "type": "record",
2041 "name": "topLevelRecord",
2042 "fields": [
2043 { "name": "id", "type": "long" },
2044 { "name": "full_name", "type": ["null","string"], "aliases": ["name"], "default": null },
2045 { "name": "is_active", "type": "boolean", "default": true }
2046 ]
2047 }"#;
2048 let mut reader = ReaderBuilder::new()
2049 .with_reader_schema(AvroSchema::new(reader_json.to_string()))
2050 .build(Cursor::new(bytes))?;
2051 let out = reader.next().unwrap()?;
2052 let full_name = out.column(1).as_string::<i32>();
2053 assert_eq!(full_name.value(0), "a");
2054 assert_eq!(full_name.value(1), "b");
2055 Ok(())
2056 }
2057
2058 #[test]
2059 fn writer_string_reader_string_null_order_second() -> Result<(), Box<dyn std::error::Error>> {
2060 let writer_schema = Schema::new(vec![Field::new("name", DataType::Utf8, false)]);
2062 let batch = RecordBatch::try_new(
2063 Arc::new(writer_schema.clone()),
2064 vec![Arc::new(StringArray::from(vec!["x", "y"])) as ArrayRef],
2065 )?;
2066 let bytes = write_ocf(&writer_schema, &[batch]);
2067
2068 let reader_json = r#"
2070 {
2071 "type":"record", "name":"topLevelRecord",
2072 "fields":[ { "name":"name", "type":["string","null"], "default":"x" } ]
2073 }"#;
2074
2075 let mut reader = ReaderBuilder::new()
2076 .with_reader_schema(AvroSchema::new(reader_json.to_string()))
2077 .build(Cursor::new(bytes))?;
2078
2079 let out = reader.next().unwrap()?;
2080 assert_eq!(out.num_rows(), 2);
2081
2082 let name = out.column(0).as_string::<i32>();
2084 assert_eq!(name.value(0), "x");
2085 assert_eq!(name.value(1), "y");
2086
2087 Ok(())
2088 }
2089
2090 #[test]
2091 fn promotion_writer_int_reader_nullable_long() -> Result<(), Box<dyn std::error::Error>> {
2092 let writer_schema = Schema::new(vec![Field::new("v", DataType::Int32, false)]);
2094 let batch = RecordBatch::try_new(
2095 Arc::new(writer_schema.clone()),
2096 vec![Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef],
2097 )?;
2098 let bytes = write_ocf(&writer_schema, &[batch]);
2099
2100 let reader_json = r#"
2102 {
2103 "type":"record", "name":"topLevelRecord",
2104 "fields":[ { "name":"v", "type":["null","long"], "default": null } ]
2105 }"#;
2106
2107 let mut reader = ReaderBuilder::new()
2108 .with_reader_schema(AvroSchema::new(reader_json.to_string()))
2109 .build(Cursor::new(bytes))?;
2110
2111 let out = reader.next().unwrap()?;
2112 assert_eq!(out.num_rows(), 3);
2113
2114 let v = out
2116 .column(0)
2117 .as_primitive::<arrow_array::types::Int64Type>();
2118 assert_eq!(v.values(), &[1, 2, 3]);
2119 assert!(
2120 out.column(0).nulls().is_none(),
2121 "expected no validity bitmap for all-valid column"
2122 );
2123
2124 Ok(())
2125 }
2126
2127 #[test]
2128 fn test_alltypes_schema_promotion_mixed() {
2129 for file in files() {
2130 let file = arrow_test_data(file);
2131 let mut promotions: HashMap<&str, &str> = HashMap::new();
2132 promotions.insert("id", "long");
2133 promotions.insert("tinyint_col", "float");
2134 promotions.insert("smallint_col", "double");
2135 promotions.insert("int_col", "double");
2136 promotions.insert("bigint_col", "double");
2137 promotions.insert("float_col", "double");
2138 promotions.insert("date_string_col", "string");
2139 promotions.insert("string_col", "string");
2140 let reader_schema = make_reader_schema_with_promotions(&file, &promotions);
2141 let batch = read_alltypes_with_reader_schema(&file, reader_schema);
2142 let expected = RecordBatch::try_from_iter_with_nullable([
2143 (
2144 "id",
2145 Arc::new(Int64Array::from(vec![4i64, 5, 6, 7, 2, 3, 0, 1])) as _,
2146 true,
2147 ),
2148 (
2149 "bool_col",
2150 Arc::new(BooleanArray::from_iter((0..8).map(|x| Some(x % 2 == 0)))) as _,
2151 true,
2152 ),
2153 (
2154 "tinyint_col",
2155 Arc::new(Float32Array::from_iter_values(
2156 (0..8).map(|x| (x % 2) as f32),
2157 )) as _,
2158 true,
2159 ),
2160 (
2161 "smallint_col",
2162 Arc::new(Float64Array::from_iter_values(
2163 (0..8).map(|x| (x % 2) as f64),
2164 )) as _,
2165 true,
2166 ),
2167 (
2168 "int_col",
2169 Arc::new(Float64Array::from_iter_values(
2170 (0..8).map(|x| (x % 2) as f64),
2171 )) as _,
2172 true,
2173 ),
2174 (
2175 "bigint_col",
2176 Arc::new(Float64Array::from_iter_values(
2177 (0..8).map(|x| ((x % 2) * 10) as f64),
2178 )) as _,
2179 true,
2180 ),
2181 (
2182 "float_col",
2183 Arc::new(Float64Array::from_iter_values(
2184 (0..8).map(|x| ((x % 2) as f32 * 1.1f32) as f64),
2185 )) as _,
2186 true,
2187 ),
2188 (
2189 "double_col",
2190 Arc::new(Float64Array::from_iter_values(
2191 (0..8).map(|x| (x % 2) as f64 * 10.1),
2192 )) as _,
2193 true,
2194 ),
2195 (
2196 "date_string_col",
2197 Arc::new(StringArray::from(vec![
2198 "03/01/09", "03/01/09", "04/01/09", "04/01/09", "02/01/09", "02/01/09",
2199 "01/01/09", "01/01/09",
2200 ])) as _,
2201 true,
2202 ),
2203 (
2204 "string_col",
2205 Arc::new(StringArray::from(
2206 (0..8)
2207 .map(|x| if x % 2 == 0 { "0" } else { "1" })
2208 .collect::<Vec<_>>(),
2209 )) as _,
2210 true,
2211 ),
2212 (
2213 "timestamp_col",
2214 Arc::new(
2215 TimestampMicrosecondArray::from_iter_values([
2216 1235865600000000, 1235865660000000, 1238544000000000, 1238544060000000, 1233446400000000, 1233446460000000, 1230768000000000, 1230768060000000, ])
2225 .with_timezone("+00:00"),
2226 ) as _,
2227 true,
2228 ),
2229 ])
2230 .unwrap();
2231 assert_eq!(batch, expected, "mismatch for file {file}");
2232 }
2233 }
2234
2235 #[test]
2236 fn test_alltypes_schema_promotion_long_to_float_only() {
2237 for file in files() {
2238 let file = arrow_test_data(file);
2239 let mut promotions: HashMap<&str, &str> = HashMap::new();
2240 promotions.insert("bigint_col", "float");
2241 let reader_schema = make_reader_schema_with_promotions(&file, &promotions);
2242 let batch = read_alltypes_with_reader_schema(&file, reader_schema);
2243 let expected = RecordBatch::try_from_iter_with_nullable([
2244 (
2245 "id",
2246 Arc::new(Int32Array::from(vec![4, 5, 6, 7, 2, 3, 0, 1])) as _,
2247 true,
2248 ),
2249 (
2250 "bool_col",
2251 Arc::new(BooleanArray::from_iter((0..8).map(|x| Some(x % 2 == 0)))) as _,
2252 true,
2253 ),
2254 (
2255 "tinyint_col",
2256 Arc::new(Int32Array::from_iter_values((0..8).map(|x| x % 2))) as _,
2257 true,
2258 ),
2259 (
2260 "smallint_col",
2261 Arc::new(Int32Array::from_iter_values((0..8).map(|x| x % 2))) as _,
2262 true,
2263 ),
2264 (
2265 "int_col",
2266 Arc::new(Int32Array::from_iter_values((0..8).map(|x| x % 2))) as _,
2267 true,
2268 ),
2269 (
2270 "bigint_col",
2271 Arc::new(Float32Array::from_iter_values(
2272 (0..8).map(|x| ((x % 2) * 10) as f32),
2273 )) as _,
2274 true,
2275 ),
2276 (
2277 "float_col",
2278 Arc::new(Float32Array::from_iter_values(
2279 (0..8).map(|x| (x % 2) as f32 * 1.1),
2280 )) as _,
2281 true,
2282 ),
2283 (
2284 "double_col",
2285 Arc::new(Float64Array::from_iter_values(
2286 (0..8).map(|x| (x % 2) as f64 * 10.1),
2287 )) as _,
2288 true,
2289 ),
2290 (
2291 "date_string_col",
2292 Arc::new(BinaryArray::from_iter_values([
2293 [48, 51, 47, 48, 49, 47, 48, 57],
2294 [48, 51, 47, 48, 49, 47, 48, 57],
2295 [48, 52, 47, 48, 49, 47, 48, 57],
2296 [48, 52, 47, 48, 49, 47, 48, 57],
2297 [48, 50, 47, 48, 49, 47, 48, 57],
2298 [48, 50, 47, 48, 49, 47, 48, 57],
2299 [48, 49, 47, 48, 49, 47, 48, 57],
2300 [48, 49, 47, 48, 49, 47, 48, 57],
2301 ])) as _,
2302 true,
2303 ),
2304 (
2305 "string_col",
2306 Arc::new(BinaryArray::from_iter_values((0..8).map(|x| [48 + x % 2]))) as _,
2307 true,
2308 ),
2309 (
2310 "timestamp_col",
2311 Arc::new(
2312 TimestampMicrosecondArray::from_iter_values([
2313 1235865600000000, 1235865660000000, 1238544000000000, 1238544060000000, 1233446400000000, 1233446460000000, 1230768000000000, 1230768060000000, ])
2322 .with_timezone("+00:00"),
2323 ) as _,
2324 true,
2325 ),
2326 ])
2327 .unwrap();
2328 assert_eq!(batch, expected, "mismatch for file {file}");
2329 }
2330 }
2331
2332 #[test]
2333 fn test_alltypes_schema_promotion_bytes_to_string_only() {
2334 for file in files() {
2335 let file = arrow_test_data(file);
2336 let mut promotions: HashMap<&str, &str> = HashMap::new();
2337 promotions.insert("date_string_col", "string");
2338 promotions.insert("string_col", "string");
2339 let reader_schema = make_reader_schema_with_promotions(&file, &promotions);
2340 let batch = read_alltypes_with_reader_schema(&file, reader_schema);
2341 let expected = RecordBatch::try_from_iter_with_nullable([
2342 (
2343 "id",
2344 Arc::new(Int32Array::from(vec![4, 5, 6, 7, 2, 3, 0, 1])) as _,
2345 true,
2346 ),
2347 (
2348 "bool_col",
2349 Arc::new(BooleanArray::from_iter((0..8).map(|x| Some(x % 2 == 0)))) as _,
2350 true,
2351 ),
2352 (
2353 "tinyint_col",
2354 Arc::new(Int32Array::from_iter_values((0..8).map(|x| x % 2))) as _,
2355 true,
2356 ),
2357 (
2358 "smallint_col",
2359 Arc::new(Int32Array::from_iter_values((0..8).map(|x| x % 2))) as _,
2360 true,
2361 ),
2362 (
2363 "int_col",
2364 Arc::new(Int32Array::from_iter_values((0..8).map(|x| x % 2))) as _,
2365 true,
2366 ),
2367 (
2368 "bigint_col",
2369 Arc::new(Int64Array::from_iter_values((0..8).map(|x| (x % 2) * 10))) as _,
2370 true,
2371 ),
2372 (
2373 "float_col",
2374 Arc::new(Float32Array::from_iter_values(
2375 (0..8).map(|x| (x % 2) as f32 * 1.1),
2376 )) as _,
2377 true,
2378 ),
2379 (
2380 "double_col",
2381 Arc::new(Float64Array::from_iter_values(
2382 (0..8).map(|x| (x % 2) as f64 * 10.1),
2383 )) as _,
2384 true,
2385 ),
2386 (
2387 "date_string_col",
2388 Arc::new(StringArray::from(vec![
2389 "03/01/09", "03/01/09", "04/01/09", "04/01/09", "02/01/09", "02/01/09",
2390 "01/01/09", "01/01/09",
2391 ])) as _,
2392 true,
2393 ),
2394 (
2395 "string_col",
2396 Arc::new(StringArray::from(
2397 (0..8)
2398 .map(|x| if x % 2 == 0 { "0" } else { "1" })
2399 .collect::<Vec<_>>(),
2400 )) as _,
2401 true,
2402 ),
2403 (
2404 "timestamp_col",
2405 Arc::new(
2406 TimestampMicrosecondArray::from_iter_values([
2407 1235865600000000, 1235865660000000, 1238544000000000, 1238544060000000, 1233446400000000, 1233446460000000, 1230768000000000, 1230768060000000, ])
2416 .with_timezone("+00:00"),
2417 ) as _,
2418 true,
2419 ),
2420 ])
2421 .unwrap();
2422 assert_eq!(batch, expected, "mismatch for file {file}");
2423 }
2424 }
2425
2426 #[test]
2427 #[cfg(feature = "snappy")]
2429 fn test_alltypes_illegal_promotion_bool_to_double_errors() {
2430 let file = arrow_test_data("avro/alltypes_plain.avro");
2431 let mut promotions: HashMap<&str, &str> = HashMap::new();
2432 promotions.insert("bool_col", "double"); let reader_schema = make_reader_schema_with_promotions(&file, &promotions);
2434 let file_handle = File::open(&file).unwrap();
2435 let result = ReaderBuilder::new()
2436 .with_reader_schema(reader_schema)
2437 .build(BufReader::new(file_handle));
2438 let err = result.expect_err("expected illegal promotion to error");
2439 let msg = err.to_string();
2440 assert!(
2441 msg.contains("Illegal promotion") || msg.contains("illegal promotion"),
2442 "unexpected error: {msg}"
2443 );
2444 }
2445
2446 #[test]
2447 fn test_simple_enum_with_reader_schema_mapping() {
2448 let file = arrow_test_data("avro/simple_enum.avro");
2449 let mut remap: HashMap<&str, Vec<&str>> = HashMap::new();
2450 remap.insert("f1", vec!["d", "c", "b", "a"]);
2451 remap.insert("f2", vec!["h", "g", "f", "e"]);
2452 remap.insert("f3", vec!["k", "i", "j"]);
2453 let reader_schema = make_reader_schema_with_enum_remap(&file, &remap);
2454 let actual = read_alltypes_with_reader_schema(&file, reader_schema);
2455 let dict_type = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
2456 let f1_keys = Int32Array::from(vec![3, 2, 1, 0]);
2458 let f1_vals = StringArray::from(vec!["d", "c", "b", "a"]);
2459 let f1 = DictionaryArray::<Int32Type>::try_new(f1_keys, Arc::new(f1_vals)).unwrap();
2460 let mut md_f1 = HashMap::new();
2461 md_f1.insert(
2462 AVRO_ENUM_SYMBOLS_METADATA_KEY.to_string(),
2463 r#"["d","c","b","a"]"#.to_string(),
2464 );
2465 md_f1.insert("avro.name".to_string(), "enum1".to_string());
2467 md_f1.insert("avro.namespace".to_string(), "ns1".to_string());
2468 let f1_field = Field::new("f1", dict_type.clone(), false).with_metadata(md_f1);
2469 let f2_keys = Int32Array::from(vec![1, 0, 3, 2]);
2471 let f2_vals = StringArray::from(vec!["h", "g", "f", "e"]);
2472 let f2 = DictionaryArray::<Int32Type>::try_new(f2_keys, Arc::new(f2_vals)).unwrap();
2473 let mut md_f2 = HashMap::new();
2474 md_f2.insert(
2475 AVRO_ENUM_SYMBOLS_METADATA_KEY.to_string(),
2476 r#"["h","g","f","e"]"#.to_string(),
2477 );
2478 md_f2.insert("avro.name".to_string(), "enum2".to_string());
2480 md_f2.insert("avro.namespace".to_string(), "ns2".to_string());
2481 let f2_field = Field::new("f2", dict_type.clone(), false).with_metadata(md_f2);
2482 let f3_keys = Int32Array::from(vec![Some(2), Some(0), None, Some(1)]);
2484 let f3_vals = StringArray::from(vec!["k", "i", "j"]);
2485 let f3 = DictionaryArray::<Int32Type>::try_new(f3_keys, Arc::new(f3_vals)).unwrap();
2486 let mut md_f3 = HashMap::new();
2487 md_f3.insert(
2488 AVRO_ENUM_SYMBOLS_METADATA_KEY.to_string(),
2489 r#"["k","i","j"]"#.to_string(),
2490 );
2491 md_f3.insert("avro.name".to_string(), "enum3".to_string());
2493 md_f3.insert("avro.namespace".to_string(), "ns1".to_string());
2494 let f3_field = Field::new("f3", dict_type.clone(), true).with_metadata(md_f3);
2495 let expected_schema = Arc::new(Schema::new(vec![f1_field, f2_field, f3_field]));
2496 let expected = RecordBatch::try_new(
2497 expected_schema,
2498 vec![Arc::new(f1) as ArrayRef, Arc::new(f2), Arc::new(f3)],
2499 )
2500 .unwrap();
2501 assert_eq!(actual, expected);
2502 }
2503
2504 #[test]
2505 fn test_schema_store_register_lookup() {
2506 let schema_int = make_record_schema(PrimitiveType::Int);
2507 let schema_long = make_record_schema(PrimitiveType::Long);
2508 let mut store = SchemaStore::new();
2509 let fp_int = store.register(schema_int.clone()).unwrap();
2510 let fp_long = store.register(schema_long.clone()).unwrap();
2511 assert_eq!(store.lookup(&fp_int).cloned(), Some(schema_int));
2512 assert_eq!(store.lookup(&fp_long).cloned(), Some(schema_long));
2513 assert_eq!(store.fingerprint_algorithm(), FingerprintAlgorithm::Rabin);
2514 }
2515
2516 #[test]
2517 fn test_unknown_fingerprint_is_error() {
2518 let (store, fp_int, _fp_long, _schema_int, schema_long) = make_two_schema_store();
2519 let unknown_fp = Fingerprint::Rabin(0xDEAD_BEEF_DEAD_BEEF);
2520 let prefix = make_prefix(unknown_fp);
2521 let mut decoder = make_decoder(&store, fp_int, &schema_long);
2522 let err = decoder.decode(&prefix).expect_err("decode should error");
2523 let msg = err.to_string();
2524 assert!(
2525 msg.contains("Unknown fingerprint"),
2526 "unexpected message: {msg}"
2527 );
2528 }
2529
2530 #[test]
2531 fn test_handle_prefix_incomplete_magic() {
2532 let (store, fp_int, _fp_long, _schema_int, schema_long) = make_two_schema_store();
2533 let mut decoder = make_decoder(&store, fp_int, &schema_long);
2534 let buf = &SINGLE_OBJECT_MAGIC[..1];
2535 let res = decoder.handle_prefix(buf).unwrap();
2536 assert_eq!(res, Some(0));
2537 assert!(decoder.pending_schema.is_none());
2538 }
2539
2540 #[test]
2541 fn test_handle_prefix_magic_mismatch() {
2542 let (store, fp_int, _fp_long, _schema_int, schema_long) = make_two_schema_store();
2543 let mut decoder = make_decoder(&store, fp_int, &schema_long);
2544 let buf = [0xFFu8, 0x00u8, 0x01u8];
2545 let res = decoder.handle_prefix(&buf).unwrap();
2546 assert!(res.is_none());
2547 }
2548
2549 #[test]
2550 fn test_handle_prefix_incomplete_fingerprint() {
2551 let (store, fp_int, fp_long, _schema_int, schema_long) = make_two_schema_store();
2552 let mut decoder = make_decoder(&store, fp_int, &schema_long);
2553 let long_bytes = match fp_long {
2554 Fingerprint::Rabin(v) => v.to_le_bytes(),
2555 Fingerprint::Id(id) => panic!("expected Rabin fingerprint, got ({id})"),
2556 Fingerprint::Id64(id) => panic!("expected Rabin fingerprint, got ({id})"),
2557 #[cfg(feature = "md5")]
2558 Fingerprint::MD5(v) => panic!("expected Rabin fingerprint, got ({v:?})"),
2559 #[cfg(feature = "sha256")]
2560 Fingerprint::SHA256(v) => panic!("expected Rabin fingerprint, got ({v:?})"),
2561 };
2562 let mut buf = Vec::from(SINGLE_OBJECT_MAGIC);
2563 buf.extend_from_slice(&long_bytes[..4]);
2564 let res = decoder.handle_prefix(&buf).unwrap();
2565 assert_eq!(res, Some(0));
2566 assert!(decoder.pending_schema.is_none());
2567 }
2568
2569 #[test]
2570 fn test_handle_prefix_valid_prefix_switches_schema() {
2571 let (store, fp_int, fp_long, _schema_int, schema_long) = make_two_schema_store();
2572 let mut decoder = make_decoder(&store, fp_int, &schema_long);
2573 let writer_schema_long = schema_long.schema().unwrap();
2574 let root_long = AvroFieldBuilder::new(&writer_schema_long).build().unwrap();
2575 let long_decoder = RecordDecoder::try_new_with_options(root_long.data_type()).unwrap();
2576 let _ = decoder.cache.insert(fp_long, long_decoder);
2577 let mut buf = Vec::from(SINGLE_OBJECT_MAGIC);
2578 match fp_long {
2579 Fingerprint::Rabin(v) => buf.extend_from_slice(&v.to_le_bytes()),
2580 Fingerprint::Id(id) => panic!("expected Rabin fingerprint, got ({id})"),
2581 Fingerprint::Id64(id) => panic!("expected Rabin fingerprint, got ({id})"),
2582 #[cfg(feature = "md5")]
2583 Fingerprint::MD5(v) => panic!("expected Rabin fingerprint, got ({v:?})"),
2584 #[cfg(feature = "sha256")]
2585 Fingerprint::SHA256(v) => panic!("expected Rabin fingerprint, got ({v:?})"),
2586 }
2587 let consumed = decoder.handle_prefix(&buf).unwrap().unwrap();
2588 assert_eq!(consumed, buf.len());
2589 assert!(decoder.pending_schema.is_some());
2590 assert_eq!(decoder.pending_schema.as_ref().unwrap().0, fp_long);
2591 }
2592
2593 #[test]
2594 fn test_decoder_projection_multiple_writer_schemas_no_reader_schema()
2595 -> Result<(), Box<dyn std::error::Error>> {
2596 let writer_v1 = AvroSchema::new(
2598 r#"{"type":"record","name":"E","fields":[{"name":"a","type":"int"},{"name":"b","type":"string"}]}"#
2599 .to_string(),
2600 );
2601 let writer_v2 = AvroSchema::new(
2602 r#"{"type":"record","name":"E","fields":[{"name":"a","type":"long"},{"name":"b","type":"string"},{"name":"c","type":"int"}]}"#
2603 .to_string(),
2604 );
2605 let mut store = SchemaStore::new();
2606 let fp1 = store.register(writer_v1)?;
2607 let fp2 = store.register(writer_v2)?;
2608 let mut decoder = ReaderBuilder::new()
2609 .with_writer_schema_store(store)
2610 .with_active_fingerprint(fp1)
2611 .with_batch_size(8)
2612 .with_projection(vec![1])
2613 .build_decoder()?;
2614 let mut msg1 = make_prefix(fp1);
2616 msg1.extend_from_slice(&encode_zigzag(1)); msg1.push((1u8) << 1);
2618 msg1.extend_from_slice(b"x");
2619 let mut msg2 = make_prefix(fp2);
2621 msg2.extend_from_slice(&encode_zigzag(2)); msg2.push((1u8) << 1);
2623 msg2.extend_from_slice(b"y");
2624 msg2.extend_from_slice(&encode_zigzag(7)); decoder.decode(&msg1)?;
2626 let batch1 = decoder.flush()?.expect("batch1");
2627 assert_eq!(batch1.num_columns(), 1);
2628 assert_eq!(batch1.schema().field(0).name(), "b");
2629 let b1 = batch1.column(0).as_string::<i32>();
2630 assert_eq!(b1.value(0), "x");
2631 decoder.decode(&msg2)?;
2632 let batch2 = decoder.flush()?.expect("batch2");
2633 assert_eq!(batch2.num_columns(), 1);
2634 assert_eq!(batch2.schema().field(0).name(), "b");
2635 let b2 = batch2.column(0).as_string::<i32>();
2636 assert_eq!(b2.value(0), "y");
2637 Ok(())
2638 }
2639
2640 #[test]
2641 fn test_two_messages_same_schema() {
2642 let writer_schema = make_value_schema(PrimitiveType::Int);
2643 let reader_schema = writer_schema.clone();
2644 let mut store = SchemaStore::new();
2645 let fp = store.register(writer_schema).unwrap();
2646 let msg1 = make_message(fp, 42);
2647 let msg2 = make_message(fp, 11);
2648 let input = [msg1.clone(), msg2.clone()].concat();
2649 let mut decoder = ReaderBuilder::new()
2650 .with_batch_size(8)
2651 .with_reader_schema(reader_schema.clone())
2652 .with_writer_schema_store(store)
2653 .with_active_fingerprint(fp)
2654 .build_decoder()
2655 .unwrap();
2656 let _ = decoder.decode(&input).unwrap();
2657 let batch = decoder.flush().unwrap().expect("batch");
2658 assert_eq!(batch.num_rows(), 2);
2659 let col = batch
2660 .column(0)
2661 .as_any()
2662 .downcast_ref::<Int32Array>()
2663 .unwrap();
2664 assert_eq!(col.value(0), 42);
2665 assert_eq!(col.value(1), 11);
2666 }
2667
2668 #[test]
2669 fn test_two_messages_schema_switch() {
2670 let w_int = make_value_schema(PrimitiveType::Int);
2671 let w_long = make_value_schema(PrimitiveType::Long);
2672 let mut store = SchemaStore::new();
2673 let fp_int = store.register(w_int).unwrap();
2674 let fp_long = store.register(w_long).unwrap();
2675 let msg_int = make_message(fp_int, 1);
2676 let msg_long = make_message(fp_long, 123456789_i64);
2677 let mut decoder = ReaderBuilder::new()
2678 .with_batch_size(8)
2679 .with_writer_schema_store(store)
2680 .with_active_fingerprint(fp_int)
2681 .build_decoder()
2682 .unwrap();
2683 let _ = decoder.decode(&msg_int).unwrap();
2684 let batch1 = decoder.flush().unwrap().expect("batch1");
2685 assert_eq!(batch1.num_rows(), 1);
2686 assert_eq!(
2687 batch1
2688 .column(0)
2689 .as_any()
2690 .downcast_ref::<Int32Array>()
2691 .unwrap()
2692 .value(0),
2693 1
2694 );
2695 let _ = decoder.decode(&msg_long).unwrap();
2696 let batch2 = decoder.flush().unwrap().expect("batch2");
2697 assert_eq!(batch2.num_rows(), 1);
2698 assert_eq!(
2699 batch2
2700 .column(0)
2701 .as_any()
2702 .downcast_ref::<Int64Array>()
2703 .unwrap()
2704 .value(0),
2705 123456789_i64
2706 );
2707 }
2708
2709 #[test]
2710 fn test_two_messages_same_schema_id() {
2711 let writer_schema = make_value_schema(PrimitiveType::Int);
2712 let reader_schema = writer_schema.clone();
2713 let id = 100u32;
2714 let mut store = SchemaStore::new_with_type(FingerprintAlgorithm::Id);
2716 let _ = store
2717 .set(Fingerprint::Id(id), writer_schema.clone())
2718 .expect("set id schema");
2719 let msg1 = make_message_id(id, 21);
2720 let msg2 = make_message_id(id, 22);
2721 let input = [msg1.clone(), msg2.clone()].concat();
2722 let mut decoder = ReaderBuilder::new()
2723 .with_batch_size(8)
2724 .with_reader_schema(reader_schema)
2725 .with_writer_schema_store(store)
2726 .with_active_fingerprint(Fingerprint::Id(id))
2727 .build_decoder()
2728 .unwrap();
2729 let _ = decoder.decode(&input).unwrap();
2730 let batch = decoder.flush().unwrap().expect("batch");
2731 assert_eq!(batch.num_rows(), 2);
2732 let col = batch
2733 .column(0)
2734 .as_any()
2735 .downcast_ref::<Int32Array>()
2736 .unwrap();
2737 assert_eq!(col.value(0), 21);
2738 assert_eq!(col.value(1), 22);
2739 }
2740
2741 #[test]
2742 fn test_unknown_id_fingerprint_is_error() {
2743 let writer_schema = make_value_schema(PrimitiveType::Int);
2744 let id_known = 7u32;
2745 let id_unknown = 9u32;
2746 let mut store = SchemaStore::new_with_type(FingerprintAlgorithm::Id);
2747 let _ = store
2748 .set(Fingerprint::Id(id_known), writer_schema.clone())
2749 .expect("set id schema");
2750 let mut decoder = ReaderBuilder::new()
2751 .with_batch_size(8)
2752 .with_reader_schema(writer_schema)
2753 .with_writer_schema_store(store)
2754 .with_active_fingerprint(Fingerprint::Id(id_known))
2755 .build_decoder()
2756 .unwrap();
2757 let prefix = make_id_prefix(id_unknown, 0);
2758 let err = decoder.decode(&prefix).expect_err("decode should error");
2759 let msg = err.to_string();
2760 assert!(
2761 msg.contains("Unknown fingerprint"),
2762 "unexpected message: {msg}"
2763 );
2764 }
2765
2766 #[test]
2767 fn test_handle_prefix_id_incomplete_magic() {
2768 let writer_schema = make_value_schema(PrimitiveType::Int);
2769 let id = 5u32;
2770 let mut store = SchemaStore::new_with_type(FingerprintAlgorithm::Id);
2771 let _ = store
2772 .set(Fingerprint::Id(id), writer_schema.clone())
2773 .expect("set id schema");
2774 let mut decoder = ReaderBuilder::new()
2775 .with_batch_size(8)
2776 .with_reader_schema(writer_schema)
2777 .with_writer_schema_store(store)
2778 .with_active_fingerprint(Fingerprint::Id(id))
2779 .build_decoder()
2780 .unwrap();
2781 let buf = &CONFLUENT_MAGIC[..0]; let res = decoder.handle_prefix(buf).unwrap();
2783 assert_eq!(res, Some(0));
2784 assert!(decoder.pending_schema.is_none());
2785 }
2786
2787 #[test]
2788 fn test_two_messages_same_schema_id64() {
2789 let writer_schema = make_value_schema(PrimitiveType::Int);
2790 let reader_schema = writer_schema.clone();
2791 let id = 100u64;
2792 let mut store = SchemaStore::new_with_type(FingerprintAlgorithm::Id64);
2794 let _ = store
2795 .set(Fingerprint::Id64(id), writer_schema.clone())
2796 .expect("set id schema");
2797 let msg1 = make_message_id64(id, 21);
2798 let msg2 = make_message_id64(id, 22);
2799 let input = [msg1.clone(), msg2.clone()].concat();
2800 let mut decoder = ReaderBuilder::new()
2801 .with_batch_size(8)
2802 .with_reader_schema(reader_schema)
2803 .with_writer_schema_store(store)
2804 .with_active_fingerprint(Fingerprint::Id64(id))
2805 .build_decoder()
2806 .unwrap();
2807 let _ = decoder.decode(&input).unwrap();
2808 let batch = decoder.flush().unwrap().expect("batch");
2809 assert_eq!(batch.num_rows(), 2);
2810 let col = batch
2811 .column(0)
2812 .as_any()
2813 .downcast_ref::<Int32Array>()
2814 .unwrap();
2815 assert_eq!(col.value(0), 21);
2816 assert_eq!(col.value(1), 22);
2817 }
2818
2819 #[test]
2820 fn test_decode_stream_with_schema() {
2821 struct TestCase<'a> {
2822 name: &'a str,
2823 schema: &'a str,
2824 expected_error: Option<&'a str>,
2825 }
2826 let tests = vec![
2827 TestCase {
2828 name: "success",
2829 schema: r#"{"type":"record","name":"test","fields":[{"name":"f2","type":"string"}]}"#,
2830 expected_error: None,
2831 },
2832 TestCase {
2833 name: "valid schema invalid data",
2834 schema: r#"{"type":"record","name":"test","fields":[{"name":"f2","type":"long"}]}"#,
2835 expected_error: Some("did not consume all bytes"),
2836 },
2837 ];
2838 for test in tests {
2839 let avro_schema = AvroSchema::new(test.schema.to_string());
2840 let mut store = SchemaStore::new();
2841 let fp = store.register(avro_schema.clone()).unwrap();
2842 let prefix = make_prefix(fp);
2843 let record_val = "some_string";
2844 let mut body = prefix;
2845 body.push((record_val.len() as u8) << 1);
2846 body.extend_from_slice(record_val.as_bytes());
2847 let decoder_res = ReaderBuilder::new()
2848 .with_batch_size(1)
2849 .with_writer_schema_store(store)
2850 .with_active_fingerprint(fp)
2851 .build_decoder();
2852 let decoder = match decoder_res {
2853 Ok(d) => d,
2854 Err(e) => {
2855 if let Some(expected) = test.expected_error {
2856 assert!(
2857 e.to_string().contains(expected),
2858 "Test '{}' failed at build – expected '{expected}', got '{e}'",
2859 test.name
2860 );
2861 continue;
2862 } else {
2863 panic!("Test '{}' failed during build: {e}", test.name);
2864 }
2865 }
2866 };
2867 let stream = Box::pin(stream::once(async { Bytes::from(body) }));
2868 let decoded_stream = decode_stream(decoder, stream);
2869 let batches_result: Result<Vec<RecordBatch>, ArrowError> =
2870 block_on(decoded_stream.try_collect());
2871 match (batches_result, test.expected_error) {
2872 (Ok(batches), None) => {
2873 let batch =
2874 arrow::compute::concat_batches(&batches[0].schema(), &batches).unwrap();
2875 let expected_field = Field::new("f2", DataType::Utf8, false);
2876 let expected_schema = Arc::new(Schema::new(vec![expected_field]));
2877 let expected_array = Arc::new(StringArray::from(vec![record_val]));
2878 let expected_batch =
2879 RecordBatch::try_new(expected_schema, vec![expected_array]).unwrap();
2880 assert_eq!(batch, expected_batch, "Test '{}'", test.name);
2881 }
2882 (Err(e), Some(expected)) => {
2883 assert!(
2884 e.to_string().contains(expected),
2885 "Test '{}' – expected error containing '{expected}', got '{e}'",
2886 test.name
2887 );
2888 }
2889 (Ok(_), Some(expected)) => {
2890 panic!(
2891 "Test '{}' expected failure ('{expected}') but succeeded",
2892 test.name
2893 );
2894 }
2895 (Err(e), None) => {
2896 panic!("Test '{}' unexpectedly failed with '{e}'", test.name);
2897 }
2898 }
2899 }
2900 }
2901
2902 #[test]
2903 fn test_utf8view_support() {
2904 struct TestHelper;
2905 impl TestHelper {
2906 fn with_utf8view(field: &Field) -> Field {
2907 match field.data_type() {
2908 DataType::Utf8 => {
2909 Field::new(field.name(), DataType::Utf8View, field.is_nullable())
2910 .with_metadata(field.metadata().clone())
2911 }
2912 _ => field.clone(),
2913 }
2914 }
2915 }
2916
2917 let field = TestHelper::with_utf8view(&Field::new("str_field", DataType::Utf8, false));
2918
2919 assert_eq!(field.data_type(), &DataType::Utf8View);
2920
2921 let array = StringViewArray::from(vec!["test1", "test2"]);
2922 let batch =
2923 RecordBatch::try_from_iter(vec![("str_field", Arc::new(array) as ArrayRef)]).unwrap();
2924
2925 assert!(batch.column(0).as_any().is::<StringViewArray>());
2926 }
2927
2928 fn make_reader_schema_with_default_fields(
2929 path: &str,
2930 default_fields: Vec<Value>,
2931 ) -> AvroSchema {
2932 let mut root = load_writer_schema_json(path);
2933 assert_eq!(root["type"], "record", "writer schema must be a record");
2934 root.as_object_mut()
2935 .expect("schema is a JSON object")
2936 .insert("fields".to_string(), Value::Array(default_fields));
2937 AvroSchema::new(root.to_string())
2938 }
2939
2940 #[test]
2941 fn test_schema_resolution_defaults_all_supported_types() {
2942 let path = "test/data/skippable_types.avro";
2943 let duration_default = "\u{0000}".repeat(12);
2944 let reader_schema = make_reader_schema_with_default_fields(
2945 path,
2946 vec![
2947 serde_json::json!({"name":"d_bool","type":"boolean","default":true}),
2948 serde_json::json!({"name":"d_int","type":"int","default":42}),
2949 serde_json::json!({"name":"d_long","type":"long","default":12345}),
2950 serde_json::json!({"name":"d_float","type":"float","default":1.5}),
2951 serde_json::json!({"name":"d_double","type":"double","default":2.25}),
2952 serde_json::json!({"name":"d_bytes","type":"bytes","default":"XYZ"}),
2953 serde_json::json!({"name":"d_string","type":"string","default":"hello"}),
2954 serde_json::json!({"name":"d_date","type":{"type":"int","logicalType":"date"},"default":0}),
2955 serde_json::json!({"name":"d_time_ms","type":{"type":"int","logicalType":"time-millis"},"default":1000}),
2956 serde_json::json!({"name":"d_time_us","type":{"type":"long","logicalType":"time-micros"},"default":2000}),
2957 serde_json::json!({"name":"d_ts_ms","type":{"type":"long","logicalType":"local-timestamp-millis"},"default":0}),
2958 serde_json::json!({"name":"d_ts_us","type":{"type":"long","logicalType":"local-timestamp-micros"},"default":0}),
2959 serde_json::json!({"name":"d_decimal","type":{"type":"bytes","logicalType":"decimal","precision":10,"scale":2},"default":""}),
2960 serde_json::json!({"name":"d_fixed","type":{"type":"fixed","name":"F4","size":4},"default":"ABCD"}),
2961 serde_json::json!({"name":"d_enum","type":{"type":"enum","name":"E","symbols":["A","B","C"]},"default":"A"}),
2962 serde_json::json!({"name":"d_duration","type":{"type":"fixed","name":"Dur","size":12,"logicalType":"duration"},"default":duration_default}),
2963 serde_json::json!({"name":"d_uuid","type":{"type":"string","logicalType":"uuid"},"default":"00000000-0000-0000-0000-000000000000"}),
2964 serde_json::json!({"name":"d_array","type":{"type":"array","items":"int"},"default":[1,2,3]}),
2965 serde_json::json!({"name":"d_map","type":{"type":"map","values":"long"},"default":{"a":1,"b":2}}),
2966 serde_json::json!({"name":"d_record","type":{
2967 "type":"record","name":"DefaultRec","fields":[
2968 {"name":"x","type":"int"},
2969 {"name":"y","type":["null","string"],"default":null}
2970 ]
2971 },"default":{"x":7}}),
2972 serde_json::json!({"name":"d_nullable_null","type":["null","int"],"default":null}),
2973 serde_json::json!({"name":"d_nullable_value","type":["int","null"],"default":123}),
2974 ],
2975 );
2976 let actual = read_alltypes_with_reader_schema(path, reader_schema);
2977 let num_rows = actual.num_rows();
2978 assert!(num_rows > 0, "skippable_types.avro should contain rows");
2979 assert_eq!(
2980 actual.num_columns(),
2981 22,
2982 "expected exactly our defaulted fields"
2983 );
2984 let mut arrays: Vec<Arc<dyn Array>> = Vec::with_capacity(22);
2985 arrays.push(Arc::new(BooleanArray::from_iter(std::iter::repeat_n(
2986 Some(true),
2987 num_rows,
2988 ))));
2989 arrays.push(Arc::new(Int32Array::from_iter_values(std::iter::repeat_n(
2990 42, num_rows,
2991 ))));
2992 arrays.push(Arc::new(Int64Array::from_iter_values(std::iter::repeat_n(
2993 12345, num_rows,
2994 ))));
2995 arrays.push(Arc::new(Float32Array::from_iter_values(
2996 std::iter::repeat_n(1.5f32, num_rows),
2997 )));
2998 arrays.push(Arc::new(Float64Array::from_iter_values(
2999 std::iter::repeat_n(2.25f64, num_rows),
3000 )));
3001 arrays.push(Arc::new(BinaryArray::from_iter_values(
3002 std::iter::repeat_n(b"XYZ".as_ref(), num_rows),
3003 )));
3004 arrays.push(Arc::new(StringArray::from_iter_values(
3005 std::iter::repeat_n("hello", num_rows),
3006 )));
3007 arrays.push(Arc::new(Date32Array::from_iter_values(
3008 std::iter::repeat_n(0, num_rows),
3009 )));
3010 arrays.push(Arc::new(Time32MillisecondArray::from_iter_values(
3011 std::iter::repeat_n(1_000, num_rows),
3012 )));
3013 arrays.push(Arc::new(Time64MicrosecondArray::from_iter_values(
3014 std::iter::repeat_n(2_000i64, num_rows),
3015 )));
3016 arrays.push(Arc::new(TimestampMillisecondArray::from_iter_values(
3017 std::iter::repeat_n(0i64, num_rows),
3018 )));
3019 arrays.push(Arc::new(TimestampMicrosecondArray::from_iter_values(
3020 std::iter::repeat_n(0i64, num_rows),
3021 )));
3022 #[cfg(feature = "small_decimals")]
3023 let decimal = Decimal64Array::from_iter_values(std::iter::repeat_n(0i64, num_rows))
3024 .with_precision_and_scale(10, 2)
3025 .unwrap();
3026 #[cfg(not(feature = "small_decimals"))]
3027 let decimal = Decimal128Array::from_iter_values(std::iter::repeat_n(0i128, num_rows))
3028 .with_precision_and_scale(10, 2)
3029 .unwrap();
3030 arrays.push(Arc::new(decimal));
3031 let fixed_iter = std::iter::repeat_n(Some(*b"ABCD"), num_rows);
3032 arrays.push(Arc::new(
3033 FixedSizeBinaryArray::try_from_sparse_iter_with_size(fixed_iter, 4).unwrap(),
3034 ));
3035 let enum_keys = Int32Array::from_iter_values(std::iter::repeat_n(0, num_rows));
3036 let enum_values = StringArray::from_iter_values(["A", "B", "C"]);
3037 let enum_arr =
3038 DictionaryArray::<Int32Type>::try_new(enum_keys, Arc::new(enum_values)).unwrap();
3039 arrays.push(Arc::new(enum_arr));
3040 let duration_values = std::iter::repeat_n(
3041 Some(IntervalMonthDayNanoType::make_value(0, 0, 0)),
3042 num_rows,
3043 );
3044 let duration_arr: IntervalMonthDayNanoArray = duration_values.collect();
3045 arrays.push(Arc::new(duration_arr));
3046 let uuid_bytes = [0u8; 16];
3047 let uuid_iter = std::iter::repeat_n(Some(uuid_bytes), num_rows);
3048 arrays.push(Arc::new(
3049 FixedSizeBinaryArray::try_from_sparse_iter_with_size(uuid_iter, 16).unwrap(),
3050 ));
3051 let item_field = Arc::new(Field::new(
3052 Field::LIST_FIELD_DEFAULT_NAME,
3053 DataType::Int32,
3054 false,
3055 ));
3056 let mut list_builder = ListBuilder::new(Int32Builder::new()).with_field(item_field);
3057 for _ in 0..num_rows {
3058 list_builder.values().append_value(1);
3059 list_builder.values().append_value(2);
3060 list_builder.values().append_value(3);
3061 list_builder.append(true);
3062 }
3063 arrays.push(Arc::new(list_builder.finish()));
3064 let values_field = Arc::new(Field::new(
3065 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
3066 DataType::Int64,
3067 false,
3068 ));
3069 let mut map_builder = MapBuilder::new(
3070 Some(builder::MapFieldNames {
3071 entry: Field::MAP_ENTRIES_FIELD_DEFAULT_NAME.to_string(),
3072 key: Field::MAP_KEY_FIELD_DEFAULT_NAME.to_string(),
3073 value: Field::MAP_VALUE_FIELD_DEFAULT_NAME.to_string(),
3074 }),
3075 StringBuilder::new(),
3076 Int64Builder::new(),
3077 )
3078 .with_values_field(values_field);
3079 for _ in 0..num_rows {
3080 let (keys, vals) = map_builder.entries();
3081 keys.append_value("a");
3082 vals.append_value(1);
3083 keys.append_value("b");
3084 vals.append_value(2);
3085 map_builder.append(true).unwrap();
3086 }
3087 arrays.push(Arc::new(map_builder.finish()));
3088 let rec_fields: Fields = Fields::from(vec![
3089 Field::new("x", DataType::Int32, false),
3090 Field::new("y", DataType::Utf8, true),
3091 ]);
3092 let mut sb = StructBuilder::new(
3093 rec_fields.clone(),
3094 vec![
3095 Box::new(Int32Builder::new()),
3096 Box::new(StringBuilder::new()),
3097 ],
3098 );
3099 for _ in 0..num_rows {
3100 sb.field_builder::<Int32Builder>(0).unwrap().append_value(7);
3101 sb.field_builder::<StringBuilder>(1).unwrap().append_null();
3102 sb.append(true);
3103 }
3104 arrays.push(Arc::new(sb.finish()));
3105 arrays.push(Arc::new(Int32Array::from_iter(std::iter::repeat_n(
3106 None::<i32>,
3107 num_rows,
3108 ))));
3109 arrays.push(Arc::new(Int32Array::from_iter_values(std::iter::repeat_n(
3110 123, num_rows,
3111 ))));
3112 let expected = RecordBatch::try_new(actual.schema(), arrays).unwrap();
3113 assert_eq!(
3114 actual, expected,
3115 "defaults should materialize correctly for all fields"
3116 );
3117 }
3118
3119 #[test]
3120 fn test_schema_resolution_default_enum_invalid_symbol_errors() {
3121 let path = "test/data/skippable_types.avro";
3122 let bad_schema = make_reader_schema_with_default_fields(
3123 path,
3124 vec![serde_json::json!({
3125 "name":"bad_enum",
3126 "type":{"type":"enum","name":"E","symbols":["A","B","C"]},
3127 "default":"Z"
3128 })],
3129 );
3130 let file = File::open(path).unwrap();
3131 let res = ReaderBuilder::new()
3132 .with_reader_schema(bad_schema)
3133 .build(BufReader::new(file));
3134 let err = res.expect_err("expected enum default validation to fail");
3135 let msg = err.to_string();
3136 let lower_msg = msg.to_lowercase();
3137 assert!(
3138 lower_msg.contains("enum")
3139 && (lower_msg.contains("symbol") || lower_msg.contains("default")),
3140 "unexpected error: {msg}"
3141 );
3142 }
3143
3144 #[test]
3145 fn test_schema_resolution_default_fixed_size_mismatch_errors() {
3146 let path = "test/data/skippable_types.avro";
3147 let bad_schema = make_reader_schema_with_default_fields(
3148 path,
3149 vec![serde_json::json!({
3150 "name":"bad_fixed",
3151 "type":{"type":"fixed","name":"F","size":4},
3152 "default":"ABC"
3153 })],
3154 );
3155 let file = File::open(path).unwrap();
3156 let res = ReaderBuilder::new()
3157 .with_reader_schema(bad_schema)
3158 .build(BufReader::new(file));
3159 let err = res.expect_err("expected fixed default validation to fail");
3160 let msg = err.to_string();
3161 let lower_msg = msg.to_lowercase();
3162 assert!(
3163 lower_msg.contains("fixed")
3164 && (lower_msg.contains("size")
3165 || lower_msg.contains("length")
3166 || lower_msg.contains("does not match")),
3167 "unexpected error: {msg}"
3168 );
3169 }
3170
3171 #[test]
3172 fn test_timestamp_with_utc_tz() {
3173 let path = arrow_test_data("avro/alltypes_plain.avro");
3174 let reader_schema =
3175 make_reader_schema_with_selected_fields_in_order(&path, &["timestamp_col"]);
3176 let file = File::open(path).unwrap();
3177 let reader = ReaderBuilder::new()
3178 .with_batch_size(1024)
3179 .with_utf8_view(false)
3180 .with_reader_schema(reader_schema)
3181 .with_tz(Tz::Utc)
3182 .build(BufReader::new(file))
3183 .unwrap();
3184 let schema = reader.schema();
3185 let batches = reader.collect::<Result<Vec<_>, _>>().unwrap();
3186 let batch = arrow::compute::concat_batches(&schema, &batches).unwrap();
3187 let expected = RecordBatch::try_from_iter_with_nullable([(
3188 "timestamp_col",
3189 Arc::new(
3190 TimestampMicrosecondArray::from_iter_values([
3191 1235865600000000, 1235865660000000, 1238544000000000, 1238544060000000, 1233446400000000, 1233446460000000, 1230768000000000, 1230768060000000, ])
3200 .with_timezone("UTC"),
3201 ) as _,
3202 true,
3203 )])
3204 .unwrap();
3205 assert_eq!(batch, expected);
3206 }
3207
3208 #[test]
3209 #[cfg(feature = "snappy")]
3211 fn test_alltypes_skip_writer_fields_keep_double_only() {
3212 let file = arrow_test_data("avro/alltypes_plain.avro");
3213 let reader_schema =
3214 make_reader_schema_with_selected_fields_in_order(&file, &["double_col"]);
3215 let batch = read_alltypes_with_reader_schema(&file, reader_schema);
3216 let expected = RecordBatch::try_from_iter_with_nullable([(
3217 "double_col",
3218 Arc::new(Float64Array::from_iter_values(
3219 (0..8).map(|x| (x % 2) as f64 * 10.1),
3220 )) as _,
3221 true,
3222 )])
3223 .unwrap();
3224 assert_eq!(batch, expected);
3225 }
3226
3227 #[test]
3228 #[cfg(feature = "snappy")]
3230 fn test_alltypes_skip_writer_fields_reorder_and_skip_many() {
3231 let file = arrow_test_data("avro/alltypes_plain.avro");
3232 let reader_schema =
3233 make_reader_schema_with_selected_fields_in_order(&file, &["timestamp_col", "id"]);
3234 let batch = read_alltypes_with_reader_schema(&file, reader_schema);
3235 let expected = RecordBatch::try_from_iter_with_nullable([
3236 (
3237 "timestamp_col",
3238 Arc::new(
3239 TimestampMicrosecondArray::from_iter_values([
3240 1235865600000000, 1235865660000000, 1238544000000000, 1238544060000000, 1233446400000000, 1233446460000000, 1230768000000000, 1230768060000000, ])
3249 .with_timezone("+00:00"),
3250 ) as _,
3251 true,
3252 ),
3253 (
3254 "id",
3255 Arc::new(Int32Array::from(vec![4, 5, 6, 7, 2, 3, 0, 1])) as _,
3256 true,
3257 ),
3258 ])
3259 .unwrap();
3260 assert_eq!(batch, expected);
3261 }
3262
3263 #[test]
3264 #[cfg_attr(miri, ignore)] fn test_skippable_types_project_each_field_individually() {
3266 let path = "test/data/skippable_types.avro";
3267 let full = read_file(path, 1024, false);
3268 let schema_full = full.schema();
3269 let num_rows = full.num_rows();
3270 let writer_json = load_writer_schema_json(path);
3271 assert_eq!(
3272 writer_json["type"], "record",
3273 "writer schema must be a record"
3274 );
3275 let fields_json = writer_json
3276 .get("fields")
3277 .and_then(|f| f.as_array())
3278 .expect("record has fields");
3279 assert_eq!(
3280 schema_full.fields().len(),
3281 fields_json.len(),
3282 "full read column count vs writer fields"
3283 );
3284 fn rebuild_list_array_with_element(
3285 col: &ArrayRef,
3286 new_elem: Arc<Field>,
3287 is_large: bool,
3288 ) -> ArrayRef {
3289 if is_large {
3290 let list = col
3291 .as_any()
3292 .downcast_ref::<LargeListArray>()
3293 .expect("expected LargeListArray");
3294 let offsets = list.offsets().clone();
3295 let values = list.values().clone();
3296 let validity = list.nulls().cloned();
3297 Arc::new(LargeListArray::try_new(new_elem, offsets, values, validity).unwrap())
3298 } else {
3299 let list = col
3300 .as_any()
3301 .downcast_ref::<ListArray>()
3302 .expect("expected ListArray");
3303 let offsets = list.offsets().clone();
3304 let values = list.values().clone();
3305 let validity = list.nulls().cloned();
3306 Arc::new(ListArray::try_new(new_elem, offsets, values, validity).unwrap())
3307 }
3308 }
3309 for (idx, f) in fields_json.iter().enumerate() {
3310 let name = f
3311 .get("name")
3312 .and_then(|n| n.as_str())
3313 .unwrap_or_else(|| panic!("field at index {idx} has no name"));
3314 let reader_schema = make_reader_schema_with_selected_fields_in_order(path, &[name]);
3315 let projected = read_alltypes_with_reader_schema(path, reader_schema);
3316 assert_eq!(
3317 projected.num_columns(),
3318 1,
3319 "projected batch should contain exactly the selected column '{name}'"
3320 );
3321 assert_eq!(
3322 projected.num_rows(),
3323 num_rows,
3324 "row count mismatch for projected column '{name}'"
3325 );
3326 let col_full = full.column(idx).clone();
3327 let full_field = schema_full.field(idx).as_ref().clone();
3328 let proj_field_ref = projected.schema().field(0).clone();
3329 let proj_field = proj_field_ref.as_ref();
3330 let top_meta = proj_field.metadata().clone();
3331 let (expected_field_ref, expected_col): (Arc<Field>, ArrayRef) =
3332 match (full_field.data_type(), proj_field.data_type()) {
3333 (&DataType::List(_), DataType::List(proj_elem)) => {
3334 let new_col =
3335 rebuild_list_array_with_element(&col_full, proj_elem.clone(), false);
3336 let nf = Field::new(
3337 full_field.name().clone(),
3338 proj_field.data_type().clone(),
3339 full_field.is_nullable(),
3340 )
3341 .with_metadata(top_meta);
3342 (Arc::new(nf), new_col)
3343 }
3344 (&DataType::LargeList(_), DataType::LargeList(proj_elem)) => {
3345 let new_col =
3346 rebuild_list_array_with_element(&col_full, proj_elem.clone(), true);
3347 let nf = Field::new(
3348 full_field.name().clone(),
3349 proj_field.data_type().clone(),
3350 full_field.is_nullable(),
3351 )
3352 .with_metadata(top_meta);
3353 (Arc::new(nf), new_col)
3354 }
3355 _ => {
3356 let nf = full_field.with_metadata(top_meta);
3357 (Arc::new(nf), col_full)
3358 }
3359 };
3360
3361 let expected = RecordBatch::try_new(
3362 Arc::new(Schema::new(vec![expected_field_ref])),
3363 vec![expected_col],
3364 )
3365 .unwrap();
3366 assert_eq!(
3367 projected, expected,
3368 "projected column '{name}' mismatch vs full read column"
3369 );
3370 }
3371 }
3372
3373 #[test]
3374 fn test_union_fields_avro_nullable_and_general_unions() {
3375 let path = "test/data/union_fields.avro";
3376 let batch = read_file(path, 1024, false);
3377 let schema = batch.schema();
3378 let idx = schema.index_of("nullable_int_nullfirst").unwrap();
3379 let a = batch.column(idx).as_primitive::<Int32Type>();
3380 assert_eq!(a.len(), 4);
3381 assert!(a.is_null(0));
3382 assert_eq!(a.value(1), 42);
3383 assert!(a.is_null(2));
3384 assert_eq!(a.value(3), 0);
3385 let idx = schema.index_of("nullable_string_nullsecond").unwrap();
3386 let s = batch
3387 .column(idx)
3388 .as_any()
3389 .downcast_ref::<StringArray>()
3390 .expect("nullable_string_nullsecond should be Utf8");
3391 assert_eq!(s.len(), 4);
3392 assert_eq!(s.value(0), "s1");
3393 assert!(s.is_null(1));
3394 assert_eq!(s.value(2), "s3");
3395 assert!(s.is_valid(3)); assert_eq!(s.value(3), "");
3397 let idx = schema.index_of("union_prim").unwrap();
3398 let u = batch
3399 .column(idx)
3400 .as_any()
3401 .downcast_ref::<UnionArray>()
3402 .expect("union_prim should be Union");
3403 let fields = match u.data_type() {
3404 DataType::Union(fields, mode) => {
3405 assert!(matches!(mode, UnionMode::Dense), "expect dense unions");
3406 fields
3407 }
3408 other => panic!("expected Union, got {other:?}"),
3409 };
3410 let tid_by_name = |name: &str| -> i8 {
3411 for (tid, f) in fields.iter() {
3412 if f.name() == name {
3413 return tid;
3414 }
3415 }
3416 panic!("union child '{name}' not found");
3417 };
3418 let expected_type_ids = vec![
3419 tid_by_name("long"),
3420 tid_by_name("int"),
3421 tid_by_name("float"),
3422 tid_by_name("double"),
3423 ];
3424 let type_ids: Vec<i8> = u.type_ids().iter().copied().collect();
3425 assert_eq!(
3426 type_ids, expected_type_ids,
3427 "branch selection for union_prim rows"
3428 );
3429 let longs = u
3430 .child(tid_by_name("long"))
3431 .as_any()
3432 .downcast_ref::<Int64Array>()
3433 .unwrap();
3434 assert_eq!(longs.len(), 1);
3435 let ints = u
3436 .child(tid_by_name("int"))
3437 .as_any()
3438 .downcast_ref::<Int32Array>()
3439 .unwrap();
3440 assert_eq!(ints.len(), 1);
3441 let floats = u
3442 .child(tid_by_name("float"))
3443 .as_any()
3444 .downcast_ref::<Float32Array>()
3445 .unwrap();
3446 assert_eq!(floats.len(), 1);
3447 let doubles = u
3448 .child(tid_by_name("double"))
3449 .as_any()
3450 .downcast_ref::<Float64Array>()
3451 .unwrap();
3452 assert_eq!(doubles.len(), 1);
3453 let idx = schema.index_of("union_bytes_vs_string").unwrap();
3454 let u = batch
3455 .column(idx)
3456 .as_any()
3457 .downcast_ref::<UnionArray>()
3458 .expect("union_bytes_vs_string should be Union");
3459 let fields = match u.data_type() {
3460 DataType::Union(fields, _) => fields,
3461 other => panic!("expected Union, got {other:?}"),
3462 };
3463 let tid_by_name = |name: &str| -> i8 {
3464 for (tid, f) in fields.iter() {
3465 if f.name() == name {
3466 return tid;
3467 }
3468 }
3469 panic!("union child '{name}' not found");
3470 };
3471 let tid_bytes = tid_by_name("bytes");
3472 let tid_string = tid_by_name("string");
3473 let type_ids: Vec<i8> = u.type_ids().iter().copied().collect();
3474 assert_eq!(
3475 type_ids,
3476 vec![tid_bytes, tid_string, tid_string, tid_bytes],
3477 "branch selection for bytes/string union"
3478 );
3479 let s_child = u
3480 .child(tid_string)
3481 .as_any()
3482 .downcast_ref::<StringArray>()
3483 .unwrap();
3484 assert_eq!(s_child.len(), 2);
3485 assert_eq!(s_child.value(0), "hello");
3486 assert_eq!(s_child.value(1), "world");
3487 let b_child = u
3488 .child(tid_bytes)
3489 .as_any()
3490 .downcast_ref::<BinaryArray>()
3491 .unwrap();
3492 assert_eq!(b_child.len(), 2);
3493 assert_eq!(b_child.value(0), &[0x00, 0xFF, 0x7F]);
3494 assert_eq!(b_child.value(1), b""); let idx = schema.index_of("union_enum_records_array_map").unwrap();
3496 let u = batch
3497 .column(idx)
3498 .as_any()
3499 .downcast_ref::<UnionArray>()
3500 .expect("union_enum_records_array_map should be Union");
3501 let fields = match u.data_type() {
3502 DataType::Union(fields, _) => fields,
3503 other => panic!("expected Union, got {other:?}"),
3504 };
3505 let mut tid_enum: Option<i8> = None;
3506 let mut tid_rec_a: Option<i8> = None;
3507 let mut tid_rec_b: Option<i8> = None;
3508 let mut tid_array: Option<i8> = None;
3509 for (tid, f) in fields.iter() {
3510 match f.data_type() {
3511 DataType::Dictionary(_, _) => tid_enum = Some(tid),
3512 DataType::Struct(childs) => {
3513 if childs.len() == 2 && childs[0].name() == "a" && childs[1].name() == "b" {
3514 tid_rec_a = Some(tid);
3515 } else if childs.len() == 2
3516 && childs[0].name() == "x"
3517 && childs[1].name() == "y"
3518 {
3519 tid_rec_b = Some(tid);
3520 }
3521 }
3522 DataType::List(_) => tid_array = Some(tid),
3523 _ => {}
3524 }
3525 }
3526 let (tid_enum, tid_rec_a, tid_rec_b, tid_array) = (
3527 tid_enum.expect("enum child"),
3528 tid_rec_a.expect("RecA child"),
3529 tid_rec_b.expect("RecB child"),
3530 tid_array.expect("array<long> child"),
3531 );
3532 let type_ids: Vec<i8> = u.type_ids().iter().copied().collect();
3533 assert_eq!(
3534 type_ids,
3535 vec![tid_enum, tid_rec_a, tid_rec_b, tid_array],
3536 "branch selection for complex union"
3537 );
3538 let dict = u
3539 .child(tid_enum)
3540 .as_any()
3541 .downcast_ref::<DictionaryArray<Int32Type>>()
3542 .unwrap();
3543 assert_eq!(dict.len(), 1);
3544 assert!(dict.is_valid(0));
3545 let rec_a = u
3546 .child(tid_rec_a)
3547 .as_any()
3548 .downcast_ref::<StructArray>()
3549 .unwrap();
3550 assert_eq!(rec_a.len(), 1);
3551 let a_val = rec_a
3552 .column_by_name("a")
3553 .unwrap()
3554 .as_any()
3555 .downcast_ref::<Int32Array>()
3556 .unwrap();
3557 assert_eq!(a_val.value(0), 7);
3558 let b_val = rec_a
3559 .column_by_name("b")
3560 .unwrap()
3561 .as_any()
3562 .downcast_ref::<StringArray>()
3563 .unwrap();
3564 assert_eq!(b_val.value(0), "x");
3565 let rec_b = u
3567 .child(tid_rec_b)
3568 .as_any()
3569 .downcast_ref::<StructArray>()
3570 .unwrap();
3571 let x_val = rec_b
3572 .column_by_name("x")
3573 .unwrap()
3574 .as_any()
3575 .downcast_ref::<Int64Array>()
3576 .unwrap();
3577 assert_eq!(x_val.value(0), 123_456_789_i64);
3578 let y_val = rec_b
3579 .column_by_name("y")
3580 .unwrap()
3581 .as_any()
3582 .downcast_ref::<BinaryArray>()
3583 .unwrap();
3584 assert_eq!(y_val.value(0), &[0xFF, 0x00]);
3585 let arr = u
3586 .child(tid_array)
3587 .as_any()
3588 .downcast_ref::<ListArray>()
3589 .unwrap();
3590 assert_eq!(arr.len(), 1);
3591 let first_values = arr.value(0);
3592 let longs = first_values.as_any().downcast_ref::<Int64Array>().unwrap();
3593 assert_eq!(longs.len(), 3);
3594 assert_eq!(longs.value(0), 1);
3595 assert_eq!(longs.value(1), 2);
3596 assert_eq!(longs.value(2), 3);
3597 let idx = schema.index_of("union_date_or_fixed4").unwrap();
3598 let u = batch
3599 .column(idx)
3600 .as_any()
3601 .downcast_ref::<UnionArray>()
3602 .expect("union_date_or_fixed4 should be Union");
3603 let fields = match u.data_type() {
3604 DataType::Union(fields, _) => fields,
3605 other => panic!("expected Union, got {other:?}"),
3606 };
3607 let mut tid_date: Option<i8> = None;
3608 let mut tid_fixed: Option<i8> = None;
3609 for (tid, f) in fields.iter() {
3610 match f.data_type() {
3611 DataType::Date32 => tid_date = Some(tid),
3612 DataType::FixedSizeBinary(4) => tid_fixed = Some(tid),
3613 _ => {}
3614 }
3615 }
3616 let (tid_date, tid_fixed) = (tid_date.expect("date"), tid_fixed.expect("fixed(4)"));
3617 let type_ids: Vec<i8> = u.type_ids().iter().copied().collect();
3618 assert_eq!(
3619 type_ids,
3620 vec![tid_date, tid_fixed, tid_date, tid_fixed],
3621 "branch selection for date/fixed4 union"
3622 );
3623 let dates = u
3624 .child(tid_date)
3625 .as_any()
3626 .downcast_ref::<Date32Array>()
3627 .unwrap();
3628 assert_eq!(dates.len(), 2);
3629 assert_eq!(dates.value(0), 19_000); assert_eq!(dates.value(1), 0); let fixed = u
3632 .child(tid_fixed)
3633 .as_any()
3634 .downcast_ref::<FixedSizeBinaryArray>()
3635 .unwrap();
3636 assert_eq!(fixed.len(), 2);
3637 assert_eq!(fixed.value(0), b"ABCD");
3638 assert_eq!(fixed.value(1), &[0x00, 0x11, 0x22, 0x33]);
3639 }
3640
3641 #[test]
3642 #[cfg_attr(miri, ignore)] fn test_union_schema_resolution_all_type_combinations() {
3644 let path = "test/data/union_fields.avro";
3645 let baseline = read_file(path, 1024, false);
3646 let baseline_schema = baseline.schema();
3647 let mut root = load_writer_schema_json(path);
3648 assert_eq!(root["type"], "record", "writer schema must be a record");
3649 let fields = root
3650 .get_mut("fields")
3651 .and_then(|f| f.as_array_mut())
3652 .expect("record has fields");
3653 fn is_named_type(obj: &Value, ty: &str, nm: &str) -> bool {
3654 obj.get("type").and_then(|v| v.as_str()) == Some(ty)
3655 && obj.get("name").and_then(|v| v.as_str()) == Some(nm)
3656 }
3657 fn is_logical(obj: &Value, prim: &str, lt: &str) -> bool {
3658 obj.get("type").and_then(|v| v.as_str()) == Some(prim)
3659 && obj.get("logicalType").and_then(|v| v.as_str()) == Some(lt)
3660 }
3661 fn find_first(arr: &[Value], pred: impl Fn(&Value) -> bool) -> Option<Value> {
3662 arr.iter().find(|v| pred(v)).cloned()
3663 }
3664 fn prim(s: &str) -> Value {
3665 Value::String(s.to_string())
3666 }
3667 for f in fields.iter_mut() {
3668 let Some(name) = f.get("name").and_then(|n| n.as_str()) else {
3669 continue;
3670 };
3671 match name {
3672 "nullable_int_nullfirst" => {
3674 f["type"] = json!(["int", "null"]);
3675 }
3676 "nullable_string_nullsecond" => {
3677 f["type"] = json!(["null", "string"]);
3678 }
3679 "union_prim" => {
3680 let orig = f["type"].as_array().unwrap().clone();
3681 let long = prim("long");
3682 let double = prim("double");
3683 let string = prim("string");
3684 let bytes = prim("bytes");
3685 let boolean = prim("boolean");
3686 assert!(orig.contains(&long));
3687 assert!(orig.contains(&double));
3688 assert!(orig.contains(&string));
3689 assert!(orig.contains(&bytes));
3690 assert!(orig.contains(&boolean));
3691 f["type"] = json!([long, double, string, bytes, boolean]);
3692 }
3693 "union_bytes_vs_string" => {
3694 f["type"] = json!(["string", "bytes"]);
3695 }
3696 "union_fixed_dur_decfix" => {
3697 let orig = f["type"].as_array().unwrap().clone();
3698 let fx8 = find_first(&orig, |o| is_named_type(o, "fixed", "Fx8")).unwrap();
3699 let dur12 = find_first(&orig, |o| is_named_type(o, "fixed", "Dur12")).unwrap();
3700 let decfix16 =
3701 find_first(&orig, |o| is_named_type(o, "fixed", "DecFix16")).unwrap();
3702 f["type"] = json!([decfix16, dur12, fx8]);
3703 }
3704 "union_enum_records_array_map" => {
3705 let orig = f["type"].as_array().unwrap().clone();
3706 let enum_color = find_first(&orig, |o| {
3707 o.get("type").and_then(|v| v.as_str()) == Some("enum")
3708 })
3709 .unwrap();
3710 let rec_a = find_first(&orig, |o| is_named_type(o, "record", "RecA")).unwrap();
3711 let rec_b = find_first(&orig, |o| is_named_type(o, "record", "RecB")).unwrap();
3712 let arr = find_first(&orig, |o| {
3713 o.get("type").and_then(|v| v.as_str()) == Some("array")
3714 })
3715 .unwrap();
3716 let map = find_first(&orig, |o| {
3717 o.get("type").and_then(|v| v.as_str()) == Some("map")
3718 })
3719 .unwrap();
3720 f["type"] = json!([arr, map, rec_b, rec_a, enum_color]);
3721 }
3722 "union_date_or_fixed4" => {
3723 let orig = f["type"].as_array().unwrap().clone();
3724 let date = find_first(&orig, |o| is_logical(o, "int", "date")).unwrap();
3725 let fx4 = find_first(&orig, |o| is_named_type(o, "fixed", "Fx4")).unwrap();
3726 f["type"] = json!([fx4, date]);
3727 }
3728 "union_time_millis_or_enum" => {
3729 let orig = f["type"].as_array().unwrap().clone();
3730 let time_ms =
3731 find_first(&orig, |o| is_logical(o, "int", "time-millis")).unwrap();
3732 let en = find_first(&orig, |o| {
3733 o.get("type").and_then(|v| v.as_str()) == Some("enum")
3734 })
3735 .unwrap();
3736 f["type"] = json!([en, time_ms]);
3737 }
3738 "union_time_micros_or_string" => {
3739 let orig = f["type"].as_array().unwrap().clone();
3740 let time_us =
3741 find_first(&orig, |o| is_logical(o, "long", "time-micros")).unwrap();
3742 f["type"] = json!(["string", time_us]);
3743 }
3744 "union_ts_millis_utc_or_array" => {
3745 let orig = f["type"].as_array().unwrap().clone();
3746 let ts_ms =
3747 find_first(&orig, |o| is_logical(o, "long", "timestamp-millis")).unwrap();
3748 let arr = find_first(&orig, |o| {
3749 o.get("type").and_then(|v| v.as_str()) == Some("array")
3750 })
3751 .unwrap();
3752 f["type"] = json!([arr, ts_ms]);
3753 }
3754 "union_ts_micros_local_or_bytes" => {
3755 let orig = f["type"].as_array().unwrap().clone();
3756 let lts_us =
3757 find_first(&orig, |o| is_logical(o, "long", "local-timestamp-micros"))
3758 .unwrap();
3759 f["type"] = json!(["bytes", lts_us]);
3760 }
3761 "union_uuid_or_fixed10" => {
3762 let orig = f["type"].as_array().unwrap().clone();
3763 let uuid = find_first(&orig, |o| is_logical(o, "string", "uuid")).unwrap();
3764 let fx10 = find_first(&orig, |o| is_named_type(o, "fixed", "Fx10")).unwrap();
3765 f["type"] = json!([fx10, uuid]);
3766 }
3767 "union_dec_bytes_or_dec_fixed" => {
3768 let orig = f["type"].as_array().unwrap().clone();
3769 let dec_bytes = find_first(&orig, |o| {
3770 o.get("type").and_then(|v| v.as_str()) == Some("bytes")
3771 && o.get("logicalType").and_then(|v| v.as_str()) == Some("decimal")
3772 })
3773 .unwrap();
3774 let dec_fix = find_first(&orig, |o| {
3775 is_named_type(o, "fixed", "DecFix20")
3776 && o.get("logicalType").and_then(|v| v.as_str()) == Some("decimal")
3777 })
3778 .unwrap();
3779 f["type"] = json!([dec_fix, dec_bytes]);
3780 }
3781 "union_null_bytes_string" => {
3782 f["type"] = json!(["bytes", "string", "null"]);
3783 }
3784 "array_of_union" => {
3785 let obj = f
3786 .get_mut("type")
3787 .expect("array type")
3788 .as_object_mut()
3789 .unwrap();
3790 obj.insert("items".to_string(), json!(["string", "long"]));
3791 }
3792 "map_of_union" => {
3793 let obj = f
3794 .get_mut("type")
3795 .expect("map type")
3796 .as_object_mut()
3797 .unwrap();
3798 obj.insert("values".to_string(), json!(["double", "null"]));
3799 }
3800 "record_with_union_field" => {
3801 let rec = f
3802 .get_mut("type")
3803 .expect("record type")
3804 .as_object_mut()
3805 .unwrap();
3806 let rec_fields = rec.get_mut("fields").unwrap().as_array_mut().unwrap();
3807 let mut found = false;
3808 for rf in rec_fields.iter_mut() {
3809 if rf.get("name").and_then(|v| v.as_str()) == Some("u") {
3810 rf["type"] = json!(["string", "long"]); found = true;
3812 break;
3813 }
3814 }
3815 assert!(found, "field 'u' expected in HasUnion");
3816 }
3817 "union_ts_micros_utc_or_map" => {
3818 let orig = f["type"].as_array().unwrap().clone();
3819 let ts_us =
3820 find_first(&orig, |o| is_logical(o, "long", "timestamp-micros")).unwrap();
3821 let map = find_first(&orig, |o| {
3822 o.get("type").and_then(|v| v.as_str()) == Some("map")
3823 })
3824 .unwrap();
3825 f["type"] = json!([map, ts_us]);
3826 }
3827 "union_ts_millis_local_or_string" => {
3828 let orig = f["type"].as_array().unwrap().clone();
3829 let lts_ms =
3830 find_first(&orig, |o| is_logical(o, "long", "local-timestamp-millis"))
3831 .unwrap();
3832 f["type"] = json!(["string", lts_ms]);
3833 }
3834 "union_bool_or_string" => {
3835 f["type"] = json!(["string", "boolean"]);
3836 }
3837 _ => {}
3838 }
3839 }
3840 let reader_schema = AvroSchema::new(root.to_string());
3841 let resolved = read_alltypes_with_reader_schema(path, reader_schema);
3842
3843 fn branch_token(dt: &DataType) -> String {
3844 match dt {
3845 DataType::Null => "null".into(),
3846 DataType::Boolean => "boolean".into(),
3847 DataType::Int32 => "int".into(),
3848 DataType::Int64 => "long".into(),
3849 DataType::Float32 => "float".into(),
3850 DataType::Float64 => "double".into(),
3851 DataType::Binary => "bytes".into(),
3852 DataType::Utf8 => "string".into(),
3853 DataType::Date32 => "date".into(),
3854 DataType::Time32(arrow_schema::TimeUnit::Millisecond) => "time-millis".into(),
3855 DataType::Time64(arrow_schema::TimeUnit::Microsecond) => "time-micros".into(),
3856 DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, tz) => if tz.is_some() {
3857 "timestamp-millis"
3858 } else {
3859 "local-timestamp-millis"
3860 }
3861 .into(),
3862 DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, tz) => if tz.is_some() {
3863 "timestamp-micros"
3864 } else {
3865 "local-timestamp-micros"
3866 }
3867 .into(),
3868 DataType::Interval(IntervalUnit::MonthDayNano) => "duration".into(),
3869 DataType::FixedSizeBinary(n) => format!("fixed{n}"),
3870 DataType::Dictionary(_, _) => "enum".into(),
3871 DataType::Decimal128(p, s) => format!("decimal({p},{s})"),
3872 DataType::Decimal256(p, s) => format!("decimal({p},{s})"),
3873 #[cfg(feature = "small_decimals")]
3874 DataType::Decimal64(p, s) => format!("decimal({p},{s})"),
3875 DataType::Struct(fields) => {
3876 if fields.len() == 2 && fields[0].name() == "a" && fields[1].name() == "b" {
3877 "record:RecA".into()
3878 } else if fields.len() == 2
3879 && fields[0].name() == "x"
3880 && fields[1].name() == "y"
3881 {
3882 "record:RecB".into()
3883 } else {
3884 "record".into()
3885 }
3886 }
3887 DataType::List(_) => "array".into(),
3888 DataType::Map(_, _) => "map".into(),
3889 other => format!("{other:?}"),
3890 }
3891 }
3892
3893 fn union_tokens(u: &UnionArray) -> (Vec<i8>, HashMap<i8, String>) {
3894 let fields = match u.data_type() {
3895 DataType::Union(fields, _) => fields,
3896 other => panic!("expected Union, got {other:?}"),
3897 };
3898 let mut dict: HashMap<i8, String> = HashMap::with_capacity(fields.len());
3899 for (tid, f) in fields.iter() {
3900 dict.insert(tid, branch_token(f.data_type()));
3901 }
3902 let ids: Vec<i8> = u.type_ids().iter().copied().collect();
3903 (ids, dict)
3904 }
3905
3906 fn expected_token(field_name: &str, writer_token: &str) -> String {
3907 match field_name {
3908 "union_prim" => match writer_token {
3909 "int" => "long".into(),
3910 "float" => "double".into(),
3911 other => other.into(),
3912 },
3913 "record_with_union_field.u" => match writer_token {
3914 "int" => "long".into(),
3915 other => other.into(),
3916 },
3917 _ => writer_token.into(),
3918 }
3919 }
3920
3921 fn get_union<'a>(
3922 rb: &'a RecordBatch,
3923 schema: arrow_schema::SchemaRef,
3924 fname: &str,
3925 ) -> &'a UnionArray {
3926 let idx = schema.index_of(fname).unwrap();
3927 rb.column(idx)
3928 .as_any()
3929 .downcast_ref::<UnionArray>()
3930 .unwrap_or_else(|| panic!("{fname} should be a Union"))
3931 }
3932
3933 fn assert_union_equivalent(field_name: &str, u_writer: &UnionArray, u_reader: &UnionArray) {
3934 let (ids_w, dict_w) = union_tokens(u_writer);
3935 let (ids_r, dict_r) = union_tokens(u_reader);
3936 assert_eq!(
3937 ids_w.len(),
3938 ids_r.len(),
3939 "{field_name}: row count mismatch between baseline and resolved"
3940 );
3941 for (i, (id_w, id_r)) in ids_w.iter().zip(ids_r.iter()).enumerate() {
3942 let w_tok = dict_w.get(id_w).unwrap();
3943 let want = expected_token(field_name, w_tok);
3944 let got = dict_r.get(id_r).unwrap();
3945 assert_eq!(
3946 got, &want,
3947 "{field_name}: row {i} resolved to wrong union branch (writer={w_tok}, expected={want}, got={got})"
3948 );
3949 }
3950 }
3951
3952 for (fname, dt) in [
3953 ("nullable_int_nullfirst", DataType::Int32),
3954 ("nullable_string_nullsecond", DataType::Utf8),
3955 ] {
3956 let idx_b = baseline_schema.index_of(fname).unwrap();
3957 let idx_r = resolved.schema().index_of(fname).unwrap();
3958 let col_b = baseline.column(idx_b);
3959 let col_r = resolved.column(idx_r);
3960 assert_eq!(
3961 col_b.data_type(),
3962 &dt,
3963 "baseline {fname} should decode as non-union with nullability"
3964 );
3965 assert_eq!(
3966 col_b.as_ref(),
3967 col_r.as_ref(),
3968 "{fname}: values must be identical regardless of null-branch order"
3969 );
3970 }
3971 let union_fields = [
3972 "union_prim",
3973 "union_bytes_vs_string",
3974 "union_fixed_dur_decfix",
3975 "union_enum_records_array_map",
3976 "union_date_or_fixed4",
3977 "union_time_millis_or_enum",
3978 "union_time_micros_or_string",
3979 "union_ts_millis_utc_or_array",
3980 "union_ts_micros_local_or_bytes",
3981 "union_uuid_or_fixed10",
3982 "union_dec_bytes_or_dec_fixed",
3983 "union_null_bytes_string",
3984 "union_ts_micros_utc_or_map",
3985 "union_ts_millis_local_or_string",
3986 "union_bool_or_string",
3987 ];
3988 for fname in union_fields {
3989 let u_b = get_union(&baseline, baseline_schema.clone(), fname);
3990 let u_r = get_union(&resolved, resolved.schema(), fname);
3991 assert_union_equivalent(fname, u_b, u_r);
3992 }
3993 {
3994 let fname = "array_of_union";
3995 let idx_b = baseline_schema.index_of(fname).unwrap();
3996 let idx_r = resolved.schema().index_of(fname).unwrap();
3997 let arr_b = baseline
3998 .column(idx_b)
3999 .as_any()
4000 .downcast_ref::<ListArray>()
4001 .expect("array_of_union should be a List");
4002 let arr_r = resolved
4003 .column(idx_r)
4004 .as_any()
4005 .downcast_ref::<ListArray>()
4006 .expect("array_of_union should be a List");
4007 assert_eq!(
4008 arr_b.value_offsets(),
4009 arr_r.value_offsets(),
4010 "{fname}: list offsets changed after resolution"
4011 );
4012 let u_b = arr_b
4013 .values()
4014 .as_any()
4015 .downcast_ref::<UnionArray>()
4016 .expect("array items should be Union");
4017 let u_r = arr_r
4018 .values()
4019 .as_any()
4020 .downcast_ref::<UnionArray>()
4021 .expect("array items should be Union");
4022 let (ids_b, dict_b) = union_tokens(u_b);
4023 let (ids_r, dict_r) = union_tokens(u_r);
4024 assert_eq!(ids_b.len(), ids_r.len(), "{fname}: values length mismatch");
4025 for (i, (id_b, id_r)) in ids_b.iter().zip(ids_r.iter()).enumerate() {
4026 let w_tok = dict_b.get(id_b).unwrap();
4027 let got = dict_r.get(id_r).unwrap();
4028 assert_eq!(
4029 got, w_tok,
4030 "{fname}: value {i} resolved to wrong branch (writer={w_tok}, got={got})"
4031 );
4032 }
4033 }
4034 {
4035 let fname = "map_of_union";
4036 let idx_b = baseline_schema.index_of(fname).unwrap();
4037 let idx_r = resolved.schema().index_of(fname).unwrap();
4038 let map_b = baseline
4039 .column(idx_b)
4040 .as_any()
4041 .downcast_ref::<MapArray>()
4042 .expect("map_of_union should be a Map");
4043 let map_r = resolved
4044 .column(idx_r)
4045 .as_any()
4046 .downcast_ref::<MapArray>()
4047 .expect("map_of_union should be a Map");
4048 assert_eq!(
4049 map_b.value_offsets(),
4050 map_r.value_offsets(),
4051 "{fname}: map value offsets changed after resolution"
4052 );
4053 let ent_b = map_b.entries();
4054 let ent_r = map_r.entries();
4055 let val_b_any = ent_b.column(1).as_ref();
4056 let val_r_any = ent_r.column(1).as_ref();
4057 let b_union = val_b_any.as_any().downcast_ref::<UnionArray>();
4058 let r_union = val_r_any.as_any().downcast_ref::<UnionArray>();
4059 if let (Some(u_b), Some(u_r)) = (b_union, r_union) {
4060 assert_union_equivalent(fname, u_b, u_r);
4061 } else {
4062 assert_eq!(
4063 val_b_any.data_type(),
4064 val_r_any.data_type(),
4065 "{fname}: value data types differ after resolution"
4066 );
4067 assert_eq!(
4068 val_b_any, val_r_any,
4069 "{fname}: value arrays differ after resolution (nullable value column case)"
4070 );
4071 let value_nullable = |m: &MapArray| -> bool {
4072 match m.data_type() {
4073 DataType::Map(entries_field, _sorted) => match entries_field.data_type() {
4074 DataType::Struct(fields) => {
4075 assert_eq!(fields.len(), 2, "entries struct must have 2 fields");
4076 assert_eq!(fields[0].name(), "key");
4077 assert_eq!(fields[1].name(), "value");
4078 fields[1].is_nullable()
4079 }
4080 other => panic!("Map entries field must be Struct, got {other:?}"),
4081 },
4082 other => panic!("expected Map data type, got {other:?}"),
4083 }
4084 };
4085 assert!(
4086 value_nullable(map_b),
4087 "{fname}: baseline Map value field should be nullable per Arrow spec"
4088 );
4089 assert!(
4090 value_nullable(map_r),
4091 "{fname}: resolved Map value field should be nullable per Arrow spec"
4092 );
4093 }
4094 }
4095 {
4096 let fname = "record_with_union_field";
4097 let idx_b = baseline_schema.index_of(fname).unwrap();
4098 let idx_r = resolved.schema().index_of(fname).unwrap();
4099 let rec_b = baseline
4100 .column(idx_b)
4101 .as_any()
4102 .downcast_ref::<StructArray>()
4103 .expect("record_with_union_field should be a Struct");
4104 let rec_r = resolved
4105 .column(idx_r)
4106 .as_any()
4107 .downcast_ref::<StructArray>()
4108 .expect("record_with_union_field should be a Struct");
4109 let u_b = rec_b
4110 .column_by_name("u")
4111 .unwrap()
4112 .as_any()
4113 .downcast_ref::<UnionArray>()
4114 .expect("field 'u' should be Union (baseline)");
4115 let u_r = rec_r
4116 .column_by_name("u")
4117 .unwrap()
4118 .as_any()
4119 .downcast_ref::<UnionArray>()
4120 .expect("field 'u' should be Union (resolved)");
4121 assert_union_equivalent("record_with_union_field.u", u_b, u_r);
4122 }
4123 }
4124
4125 #[test]
4126 fn test_union_fields_end_to_end_expected_arrays() {
4127 fn tid_by_name(fields: &UnionFields, want: &str) -> i8 {
4128 for (tid, f) in fields.iter() {
4129 if f.name() == want {
4130 return tid;
4131 }
4132 }
4133 panic!("union child '{want}' not found")
4134 }
4135
4136 fn tid_by_dt(fields: &UnionFields, pred: impl Fn(&DataType) -> bool) -> i8 {
4137 for (tid, f) in fields.iter() {
4138 if pred(f.data_type()) {
4139 return tid;
4140 }
4141 }
4142 panic!("no union child matches predicate");
4143 }
4144
4145 fn uuid16_from_str(s: &str) -> [u8; 16] {
4146 fn hex(b: u8) -> u8 {
4147 match b {
4148 b'0'..=b'9' => b - b'0',
4149 b'a'..=b'f' => b - b'a' + 10,
4150 b'A'..=b'F' => b - b'A' + 10,
4151 _ => panic!("invalid hex"),
4152 }
4153 }
4154 let mut out = [0u8; 16];
4155 let bytes = s.as_bytes();
4156 let (mut i, mut j) = (0, 0);
4157 while i < bytes.len() {
4158 if bytes[i] == b'-' {
4159 i += 1;
4160 continue;
4161 }
4162 let hi = hex(bytes[i]);
4163 let lo = hex(bytes[i + 1]);
4164 out[j] = (hi << 4) | lo;
4165 j += 1;
4166 i += 2;
4167 }
4168 assert_eq!(j, 16, "uuid must decode to 16 bytes");
4169 out
4170 }
4171
4172 fn empty_child_for(dt: &DataType) -> Arc<dyn Array> {
4173 match dt {
4174 DataType::Null => Arc::new(NullArray::new(0)),
4175 DataType::Boolean => Arc::new(BooleanArray::from(Vec::<bool>::new())),
4176 DataType::Int32 => Arc::new(Int32Array::from(Vec::<i32>::new())),
4177 DataType::Int64 => Arc::new(Int64Array::from(Vec::<i64>::new())),
4178 DataType::Float32 => Arc::new(arrow_array::Float32Array::from(Vec::<f32>::new())),
4179 DataType::Float64 => Arc::new(arrow_array::Float64Array::from(Vec::<f64>::new())),
4180 DataType::Binary => Arc::new(BinaryArray::from(Vec::<&[u8]>::new())),
4181 DataType::Utf8 => Arc::new(StringArray::from(Vec::<&str>::new())),
4182 DataType::Date32 => Arc::new(arrow_array::Date32Array::from(Vec::<i32>::new())),
4183 DataType::Time32(arrow_schema::TimeUnit::Millisecond) => {
4184 Arc::new(Time32MillisecondArray::from(Vec::<i32>::new()))
4185 }
4186 DataType::Time64(arrow_schema::TimeUnit::Microsecond) => {
4187 Arc::new(Time64MicrosecondArray::from(Vec::<i64>::new()))
4188 }
4189 DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, tz) => {
4190 let a = TimestampMillisecondArray::from(Vec::<i64>::new());
4191 Arc::new(if let Some(tz) = tz {
4192 a.with_timezone(tz.clone())
4193 } else {
4194 a
4195 })
4196 }
4197 DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, tz) => {
4198 let a = TimestampMicrosecondArray::from(Vec::<i64>::new());
4199 Arc::new(if let Some(tz) = tz {
4200 a.with_timezone(tz.clone())
4201 } else {
4202 a
4203 })
4204 }
4205 DataType::Interval(IntervalUnit::MonthDayNano) => {
4206 Arc::new(arrow_array::IntervalMonthDayNanoArray::from(Vec::<
4207 IntervalMonthDayNano,
4208 >::new(
4209 )))
4210 }
4211 DataType::FixedSizeBinary(n) => Arc::new(FixedSizeBinaryArray::new_null(*n, 0)),
4212 DataType::Dictionary(k, v) => {
4213 assert_eq!(**k, DataType::Int32, "expect int32 keys for enums");
4214 let keys = Int32Array::from(Vec::<i32>::new());
4215 let values = match v.as_ref() {
4216 DataType::Utf8 => {
4217 Arc::new(StringArray::from(Vec::<&str>::new())) as ArrayRef
4218 }
4219 other => panic!("unexpected dictionary value type {other:?}"),
4220 };
4221 Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap())
4222 }
4223 DataType::List(field) => {
4224 let values: ArrayRef = match field.data_type() {
4225 DataType::Int32 => {
4226 Arc::new(Int32Array::from(Vec::<i32>::new())) as ArrayRef
4227 }
4228 DataType::Int64 => {
4229 Arc::new(Int64Array::from(Vec::<i64>::new())) as ArrayRef
4230 }
4231 DataType::Utf8 => {
4232 Arc::new(StringArray::from(Vec::<&str>::new())) as ArrayRef
4233 }
4234 DataType::Union(_, _) => {
4235 let (uf, _) = if let DataType::Union(f, m) = field.data_type() {
4236 (f.clone(), m)
4237 } else {
4238 unreachable!()
4239 };
4240 let children: Vec<ArrayRef> = uf
4241 .iter()
4242 .map(|(_, f)| empty_child_for(f.data_type()))
4243 .collect();
4244 Arc::new(
4245 UnionArray::try_new(
4246 uf.clone(),
4247 ScalarBuffer::<i8>::from(Vec::<i8>::new()),
4248 Some(ScalarBuffer::<i32>::from(Vec::<i32>::new())),
4249 children,
4250 )
4251 .unwrap(),
4252 ) as ArrayRef
4253 }
4254 other => panic!("unsupported list item type: {other:?}"),
4255 };
4256 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0]));
4257 Arc::new(ListArray::try_new(field.clone(), offsets, values, None).unwrap())
4258 }
4259 DataType::Map(entry_field, ordered) => {
4260 let DataType::Struct(childs) = entry_field.data_type() else {
4261 panic!("map entries must be struct")
4262 };
4263 let key_field = &childs[0];
4264 let val_field = &childs[1];
4265 assert_eq!(key_field.data_type(), &DataType::Utf8);
4266 let keys = StringArray::from(Vec::<&str>::new());
4267 let vals: ArrayRef = match val_field.data_type() {
4268 DataType::Float64 => {
4269 Arc::new(arrow_array::Float64Array::from(Vec::<f64>::new())) as ArrayRef
4270 }
4271 DataType::Int64 => {
4272 Arc::new(Int64Array::from(Vec::<i64>::new())) as ArrayRef
4273 }
4274 DataType::Utf8 => {
4275 Arc::new(StringArray::from(Vec::<&str>::new())) as ArrayRef
4276 }
4277 DataType::Union(uf, _) => {
4278 let ch: Vec<ArrayRef> = uf
4279 .iter()
4280 .map(|(_, f)| empty_child_for(f.data_type()))
4281 .collect();
4282 Arc::new(
4283 UnionArray::try_new(
4284 uf.clone(),
4285 ScalarBuffer::<i8>::from(Vec::<i8>::new()),
4286 Some(ScalarBuffer::<i32>::from(Vec::<i32>::new())),
4287 ch,
4288 )
4289 .unwrap(),
4290 ) as ArrayRef
4291 }
4292 other => panic!("unsupported map value type: {other:?}"),
4293 };
4294 let entries = StructArray::new(
4295 Fields::from(vec![key_field.as_ref().clone(), val_field.as_ref().clone()]),
4296 vec![Arc::new(keys) as ArrayRef, vals],
4297 None,
4298 );
4299 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0]));
4300 Arc::new(MapArray::new(
4301 entry_field.clone(),
4302 offsets,
4303 entries,
4304 None,
4305 *ordered,
4306 ))
4307 }
4308 other => panic!("empty_child_for: unhandled type {other:?}"),
4309 }
4310 }
4311
4312 fn mk_dense_union(
4313 fields: &UnionFields,
4314 type_ids: Vec<i8>,
4315 offsets: Vec<i32>,
4316 provide: impl Fn(&Field) -> Option<ArrayRef>,
4317 ) -> ArrayRef {
4318 let children: Vec<ArrayRef> = fields
4319 .iter()
4320 .map(|(_, f)| provide(f).unwrap_or_else(|| empty_child_for(f.data_type())))
4321 .collect();
4322
4323 Arc::new(
4324 UnionArray::try_new(
4325 fields.clone(),
4326 ScalarBuffer::<i8>::from(type_ids),
4327 Some(ScalarBuffer::<i32>::from(offsets)),
4328 children,
4329 )
4330 .unwrap(),
4331 ) as ArrayRef
4332 }
4333
4334 let date_a: i32 = 19_000;
4336 let time_ms_a: i32 = 13 * 3_600_000 + 45 * 60_000 + 30_000 + 123;
4337 let time_us_b: i64 = 23 * 3_600_000_000 + 59 * 60_000_000 + 59 * 1_000_000 + 999_999;
4338 let ts_ms_2024_01_01: i64 = 1_704_067_200_000;
4339 let ts_us_2024_01_01: i64 = ts_ms_2024_01_01 * 1000;
4340 let fx8_a: [u8; 8] = *b"ABCDEFGH";
4342 let fx4_abcd: [u8; 4] = *b"ABCD";
4343 let fx4_misc: [u8; 4] = [0x00, 0x11, 0x22, 0x33];
4344 let fx10_ascii: [u8; 10] = *b"0123456789";
4345 let fx10_aa: [u8; 10] = [0xAA; 10];
4346 let dur_a = IntervalMonthDayNanoType::make_value(1, 2, 3_000_000_000);
4348 let dur_b = IntervalMonthDayNanoType::make_value(12, 31, 999_000_000);
4349 let uuid1 = uuid16_from_str("fe7bc30b-4ce8-4c5e-b67c-2234a2d38e66");
4351 let uuid2 = uuid16_from_str("0826cc06-d2e3-4599-b4ad-af5fa6905cdb");
4352 let dec_b_scale2_pos: i128 = 123_456; let dec_fix16_neg: i128 = -101; let dec_fix20_s4: i128 = 1_234_567_891_234; let dec_fix20_s4_neg: i128 = -123; let path = "test/data/union_fields.avro";
4358 let actual = read_file(path, 1024, false);
4359 let schema = actual.schema();
4360 let get_union = |name: &str| -> (UnionFields, UnionMode) {
4362 let idx = schema.index_of(name).unwrap();
4363 match schema.field(idx).data_type() {
4364 DataType::Union(f, m) => (f.clone(), *m),
4365 other => panic!("{name} should be a Union, got {other:?}"),
4366 }
4367 };
4368 let mut expected_cols: Vec<ArrayRef> = Vec::with_capacity(schema.fields().len());
4369 expected_cols.push(Arc::new(Int32Array::from(vec![
4371 None,
4372 Some(42),
4373 None,
4374 Some(0),
4375 ])));
4376 expected_cols.push(Arc::new(StringArray::from(vec![
4378 Some("s1"),
4379 None,
4380 Some("s3"),
4381 Some(""),
4382 ])));
4383 {
4385 let (uf, mode) = get_union("union_prim");
4386 assert!(matches!(mode, UnionMode::Dense));
4387 let generated_names: Vec<&str> = uf.iter().map(|(_, f)| f.name().as_str()).collect();
4388 let expected_names = vec![
4389 "boolean", "int", "long", "float", "double", "bytes", "string",
4390 ];
4391 assert_eq!(
4392 generated_names, expected_names,
4393 "Field names for union_prim are incorrect"
4394 );
4395 let tids = vec![
4396 tid_by_name(&uf, "long"),
4397 tid_by_name(&uf, "int"),
4398 tid_by_name(&uf, "float"),
4399 tid_by_name(&uf, "double"),
4400 ];
4401 let offs = vec![0, 0, 0, 0];
4402 let arr = mk_dense_union(&uf, tids, offs, |f| match f.name().as_str() {
4403 "int" => Some(Arc::new(Int32Array::from(vec![-1])) as ArrayRef),
4404 "long" => Some(Arc::new(Int64Array::from(vec![1_234_567_890_123i64])) as ArrayRef),
4405 "float" => {
4406 Some(Arc::new(arrow_array::Float32Array::from(vec![1.25f32])) as ArrayRef)
4407 }
4408 "double" => {
4409 Some(Arc::new(arrow_array::Float64Array::from(vec![-2.5f64])) as ArrayRef)
4410 }
4411 _ => None,
4412 });
4413 expected_cols.push(arr);
4414 }
4415 {
4417 let (uf, _) = get_union("union_bytes_vs_string");
4418 let tids = vec![
4419 tid_by_name(&uf, "bytes"),
4420 tid_by_name(&uf, "string"),
4421 tid_by_name(&uf, "string"),
4422 tid_by_name(&uf, "bytes"),
4423 ];
4424 let offs = vec![0, 0, 1, 1];
4425 let arr = mk_dense_union(&uf, tids, offs, |f| match f.name().as_str() {
4426 "bytes" => Some(
4427 Arc::new(BinaryArray::from(vec![&[0x00, 0xFF, 0x7F][..], &[][..]])) as ArrayRef,
4428 ),
4429 "string" => Some(Arc::new(StringArray::from(vec!["hello", "world"])) as ArrayRef),
4430 _ => None,
4431 });
4432 expected_cols.push(arr);
4433 }
4434 {
4436 let (uf, _) = get_union("union_fixed_dur_decfix");
4437 let tid_fx8 = tid_by_dt(&uf, |dt| matches!(dt, DataType::FixedSizeBinary(8)));
4438 let tid_dur = tid_by_dt(&uf, |dt| {
4439 matches!(
4440 dt,
4441 DataType::Interval(arrow_schema::IntervalUnit::MonthDayNano)
4442 )
4443 });
4444 let tid_dec = tid_by_dt(&uf, |dt| match dt {
4445 #[cfg(feature = "small_decimals")]
4446 DataType::Decimal64(10, 2) => true,
4447 DataType::Decimal128(10, 2) | DataType::Decimal256(10, 2) => true,
4448 _ => false,
4449 });
4450 let tids = vec![tid_fx8, tid_dur, tid_dec, tid_dur];
4451 let offs = vec![0, 0, 0, 1];
4452 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
4453 DataType::FixedSizeBinary(8) => {
4454 let it = std::iter::once(Some(fx8_a));
4455 Some(Arc::new(
4456 FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 8).unwrap(),
4457 ) as ArrayRef)
4458 }
4459 DataType::Interval(IntervalUnit::MonthDayNano) => {
4460 Some(Arc::new(arrow_array::IntervalMonthDayNanoArray::from(vec![
4461 dur_a, dur_b,
4462 ])) as ArrayRef)
4463 }
4464 #[cfg(feature = "small_decimals")]
4465 DataType::Decimal64(10, 2) => {
4466 let a = arrow_array::Decimal64Array::from_iter_values([dec_fix16_neg as i64]);
4467 Some(Arc::new(a.with_precision_and_scale(10, 2).unwrap()) as ArrayRef)
4468 }
4469 DataType::Decimal128(10, 2) => {
4470 let a = arrow_array::Decimal128Array::from_iter_values([dec_fix16_neg]);
4471 Some(Arc::new(a.with_precision_and_scale(10, 2).unwrap()) as ArrayRef)
4472 }
4473 DataType::Decimal256(10, 2) => {
4474 let a = arrow_array::Decimal256Array::from_iter_values([i256::from_i128(
4475 dec_fix16_neg,
4476 )]);
4477 Some(Arc::new(a.with_precision_and_scale(10, 2).unwrap()) as ArrayRef)
4478 }
4479 _ => None,
4480 });
4481 let generated_names: Vec<&str> = uf.iter().map(|(_, f)| f.name().as_str()).collect();
4482 let expected_names = vec!["Fx8", "Dur12", "DecFix16"];
4483 assert_eq!(
4484 generated_names, expected_names,
4485 "Data type names were not generated correctly for union_fixed_dur_decfix"
4486 );
4487 expected_cols.push(arr);
4488 }
4489 {
4491 let (uf, _) = get_union("union_enum_records_array_map");
4492 let tid_enum = tid_by_dt(&uf, |dt| matches!(dt, DataType::Dictionary(_, _)));
4493 let tid_reca = tid_by_dt(&uf, |dt| {
4494 if let DataType::Struct(fs) = dt {
4495 fs.len() == 2 && fs[0].name() == "a" && fs[1].name() == "b"
4496 } else {
4497 false
4498 }
4499 });
4500 let tid_recb = tid_by_dt(&uf, |dt| {
4501 if let DataType::Struct(fs) = dt {
4502 fs.len() == 2 && fs[0].name() == "x" && fs[1].name() == "y"
4503 } else {
4504 false
4505 }
4506 });
4507 let tid_arr = tid_by_dt(&uf, |dt| matches!(dt, DataType::List(_)));
4508 let tids = vec![tid_enum, tid_reca, tid_recb, tid_arr];
4509 let offs = vec![0, 0, 0, 0];
4510 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
4511 DataType::Dictionary(_, _) => {
4512 let keys = Int32Array::from(vec![0i32]); let values =
4514 Arc::new(StringArray::from(vec!["RED", "GREEN", "BLUE"])) as ArrayRef;
4515 Some(
4516 Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap())
4517 as ArrayRef,
4518 )
4519 }
4520 DataType::Struct(fs)
4521 if fs.len() == 2 && fs[0].name() == "a" && fs[1].name() == "b" =>
4522 {
4523 let a = Int32Array::from(vec![7]);
4524 let b = StringArray::from(vec!["x"]);
4525 Some(Arc::new(StructArray::new(
4526 fs.clone(),
4527 vec![Arc::new(a), Arc::new(b)],
4528 None,
4529 )) as ArrayRef)
4530 }
4531 DataType::Struct(fs)
4532 if fs.len() == 2 && fs[0].name() == "x" && fs[1].name() == "y" =>
4533 {
4534 let x = Int64Array::from(vec![123_456_789i64]);
4535 let y = BinaryArray::from(vec![&[0xFF, 0x00][..]]);
4536 Some(Arc::new(StructArray::new(
4537 fs.clone(),
4538 vec![Arc::new(x), Arc::new(y)],
4539 None,
4540 )) as ArrayRef)
4541 }
4542 DataType::List(field) => {
4543 let values = Int64Array::from(vec![1i64, 2, 3]);
4544 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3]));
4545 Some(Arc::new(
4546 ListArray::try_new(field.clone(), offsets, Arc::new(values), None).unwrap(),
4547 ) as ArrayRef)
4548 }
4549 DataType::Map(_, _) => None,
4550 other => panic!("unexpected child {other:?}"),
4551 });
4552 expected_cols.push(arr);
4553 }
4554 {
4556 let (uf, _) = get_union("union_date_or_fixed4");
4557 let tid_date = tid_by_dt(&uf, |dt| matches!(dt, DataType::Date32));
4558 let tid_fx4 = tid_by_dt(&uf, |dt| matches!(dt, DataType::FixedSizeBinary(4)));
4559 let tids = vec![tid_date, tid_fx4, tid_date, tid_fx4];
4560 let offs = vec![0, 0, 1, 1];
4561 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
4562 DataType::Date32 => {
4563 Some(Arc::new(arrow_array::Date32Array::from(vec![date_a, 0])) as ArrayRef)
4564 }
4565 DataType::FixedSizeBinary(4) => {
4566 let it = [Some(fx4_abcd), Some(fx4_misc)].into_iter();
4567 Some(Arc::new(
4568 FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 4).unwrap(),
4569 ) as ArrayRef)
4570 }
4571 _ => None,
4572 });
4573 expected_cols.push(arr);
4574 }
4575 {
4577 let (uf, _) = get_union("union_time_millis_or_enum");
4578 let tid_ms = tid_by_dt(&uf, |dt| {
4579 matches!(dt, DataType::Time32(arrow_schema::TimeUnit::Millisecond))
4580 });
4581 let tid_en = tid_by_dt(&uf, |dt| matches!(dt, DataType::Dictionary(_, _)));
4582 let tids = vec![tid_ms, tid_en, tid_en, tid_ms];
4583 let offs = vec![0, 0, 1, 1];
4584 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
4585 DataType::Time32(arrow_schema::TimeUnit::Millisecond) => {
4586 Some(Arc::new(Time32MillisecondArray::from(vec![time_ms_a, 0])) as ArrayRef)
4587 }
4588 DataType::Dictionary(_, _) => {
4589 let keys = Int32Array::from(vec![0i32, 1]); let values = Arc::new(StringArray::from(vec!["ON", "OFF"])) as ArrayRef;
4591 Some(
4592 Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap())
4593 as ArrayRef,
4594 )
4595 }
4596 _ => None,
4597 });
4598 expected_cols.push(arr);
4599 }
4600 {
4602 let (uf, _) = get_union("union_time_micros_or_string");
4603 let tid_us = tid_by_dt(&uf, |dt| {
4604 matches!(dt, DataType::Time64(arrow_schema::TimeUnit::Microsecond))
4605 });
4606 let tid_s = tid_by_name(&uf, "string");
4607 let tids = vec![tid_s, tid_us, tid_s, tid_s];
4608 let offs = vec![0, 0, 1, 2];
4609 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
4610 DataType::Time64(arrow_schema::TimeUnit::Microsecond) => {
4611 Some(Arc::new(Time64MicrosecondArray::from(vec![time_us_b])) as ArrayRef)
4612 }
4613 DataType::Utf8 => {
4614 Some(Arc::new(StringArray::from(vec!["evening", "night", ""])) as ArrayRef)
4615 }
4616 _ => None,
4617 });
4618 expected_cols.push(arr);
4619 }
4620 {
4622 let (uf, _) = get_union("union_ts_millis_utc_or_array");
4623 let tid_ts = tid_by_dt(&uf, |dt| {
4624 matches!(
4625 dt,
4626 DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, _)
4627 )
4628 });
4629 let tid_arr = tid_by_dt(&uf, |dt| matches!(dt, DataType::List(_)));
4630 let tids = vec![tid_ts, tid_arr, tid_arr, tid_ts];
4631 let offs = vec![0, 0, 1, 1];
4632 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
4633 DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, tz) => {
4634 let a = TimestampMillisecondArray::from(vec![
4635 ts_ms_2024_01_01,
4636 ts_ms_2024_01_01 + 86_400_000,
4637 ]);
4638 Some(Arc::new(if let Some(tz) = tz {
4639 a.with_timezone(tz.clone())
4640 } else {
4641 a
4642 }) as ArrayRef)
4643 }
4644 DataType::List(field) => {
4645 let values = Int32Array::from(vec![0, 1, 2, -1, 0, 1]);
4646 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 6]));
4647 Some(Arc::new(
4648 ListArray::try_new(field.clone(), offsets, Arc::new(values), None).unwrap(),
4649 ) as ArrayRef)
4650 }
4651 _ => None,
4652 });
4653 expected_cols.push(arr);
4654 }
4655 {
4657 let (uf, _) = get_union("union_ts_micros_local_or_bytes");
4658 let tid_lts = tid_by_dt(&uf, |dt| {
4659 matches!(
4660 dt,
4661 DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, None)
4662 )
4663 });
4664 let tid_b = tid_by_name(&uf, "bytes");
4665 let tids = vec![tid_b, tid_lts, tid_b, tid_b];
4666 let offs = vec![0, 0, 1, 2];
4667 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
4668 DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, None) => Some(Arc::new(
4669 TimestampMicrosecondArray::from(vec![ts_us_2024_01_01]),
4670 )
4671 as ArrayRef),
4672 DataType::Binary => Some(Arc::new(BinaryArray::from(vec![
4673 &b"\x11\x22\x33"[..],
4674 &b"\x00"[..],
4675 &b"\x10\x20\x30\x40"[..],
4676 ])) as ArrayRef),
4677 _ => None,
4678 });
4679 expected_cols.push(arr);
4680 }
4681 {
4683 let (uf, _) = get_union("union_uuid_or_fixed10");
4684 let tid_fx16 = tid_by_dt(&uf, |dt| matches!(dt, DataType::FixedSizeBinary(16)));
4685 let tid_fx10 = tid_by_dt(&uf, |dt| matches!(dt, DataType::FixedSizeBinary(10)));
4686 let tids = vec![tid_fx16, tid_fx10, tid_fx16, tid_fx10];
4687 let offs = vec![0, 0, 1, 1];
4688 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
4689 DataType::FixedSizeBinary(16) => {
4690 let it = [Some(uuid1), Some(uuid2)].into_iter();
4691 Some(Arc::new(
4692 FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 16).unwrap(),
4693 ) as ArrayRef)
4694 }
4695 DataType::FixedSizeBinary(10) => {
4696 let it = [Some(fx10_ascii), Some(fx10_aa)].into_iter();
4697 Some(Arc::new(
4698 FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 10).unwrap(),
4699 ) as ArrayRef)
4700 }
4701 _ => None,
4702 });
4703 expected_cols.push(arr);
4704 }
4705 {
4707 let (uf, _) = get_union("union_dec_bytes_or_dec_fixed");
4708 let tid_b10s2 = tid_by_dt(&uf, |dt| match dt {
4709 #[cfg(feature = "small_decimals")]
4710 DataType::Decimal64(10, 2) => true,
4711 DataType::Decimal128(10, 2) | DataType::Decimal256(10, 2) => true,
4712 _ => false,
4713 });
4714 let tid_f20s4 = tid_by_dt(&uf, |dt| {
4715 matches!(
4716 dt,
4717 DataType::Decimal128(20, 4) | DataType::Decimal256(20, 4)
4718 )
4719 });
4720 let tids = vec![tid_b10s2, tid_f20s4, tid_b10s2, tid_f20s4];
4721 let offs = vec![0, 0, 1, 1];
4722 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
4723 #[cfg(feature = "small_decimals")]
4724 DataType::Decimal64(10, 2) => {
4725 let a = Decimal64Array::from_iter_values([dec_b_scale2_pos as i64, 0i64]);
4726 Some(Arc::new(a.with_precision_and_scale(10, 2).unwrap()) as ArrayRef)
4727 }
4728 DataType::Decimal128(10, 2) => {
4729 let a = Decimal128Array::from_iter_values([dec_b_scale2_pos, 0]);
4730 Some(Arc::new(a.with_precision_and_scale(10, 2).unwrap()) as ArrayRef)
4731 }
4732 DataType::Decimal256(10, 2) => {
4733 let a = Decimal256Array::from_iter_values([
4734 i256::from_i128(dec_b_scale2_pos),
4735 i256::from(0),
4736 ]);
4737 Some(Arc::new(a.with_precision_and_scale(10, 2).unwrap()) as ArrayRef)
4738 }
4739 DataType::Decimal128(20, 4) => {
4740 let a = Decimal128Array::from_iter_values([dec_fix20_s4_neg, dec_fix20_s4]);
4741 Some(Arc::new(a.with_precision_and_scale(20, 4).unwrap()) as ArrayRef)
4742 }
4743 DataType::Decimal256(20, 4) => {
4744 let a = Decimal256Array::from_iter_values([
4745 i256::from_i128(dec_fix20_s4_neg),
4746 i256::from_i128(dec_fix20_s4),
4747 ]);
4748 Some(Arc::new(a.with_precision_and_scale(20, 4).unwrap()) as ArrayRef)
4749 }
4750 _ => None,
4751 });
4752 expected_cols.push(arr);
4753 }
4754 {
4756 let (uf, _) = get_union("union_null_bytes_string");
4757 let tid_n = tid_by_name(&uf, "null");
4758 let tid_b = tid_by_name(&uf, "bytes");
4759 let tid_s = tid_by_name(&uf, "string");
4760 let tids = vec![tid_n, tid_b, tid_s, tid_s];
4761 let offs = vec![0, 0, 0, 1];
4762 let arr = mk_dense_union(&uf, tids, offs, |f| match f.name().as_str() {
4763 "null" => Some(Arc::new(arrow_array::NullArray::new(1)) as ArrayRef),
4764 "bytes" => Some(Arc::new(BinaryArray::from(vec![&b"\x01\x02"[..]])) as ArrayRef),
4765 "string" => Some(Arc::new(StringArray::from(vec!["text", "u"])) as ArrayRef),
4766 _ => None,
4767 });
4768 expected_cols.push(arr);
4769 }
4770 {
4772 let idx = schema.index_of("array_of_union").unwrap();
4773 let dt = schema.field(idx).data_type().clone();
4774 let item_field = match &dt {
4775 DataType::List(f) => f.clone(),
4776 other => panic!("array_of_union must be List, got {other:?}"),
4777 };
4778 let (uf, _) = match item_field.data_type() {
4779 DataType::Union(f, m) => (f.clone(), m),
4780 other => panic!("array_of_union items must be Union, got {other:?}"),
4781 };
4782 let tid_l = tid_by_name(&uf, "long");
4783 let tid_s = tid_by_name(&uf, "string");
4784 let type_ids = vec![tid_l, tid_s, tid_l, tid_s, tid_l, tid_l, tid_s, tid_l];
4785 let offsets = vec![0, 0, 1, 1, 2, 3, 2, 4];
4786 let values_union =
4787 mk_dense_union(&uf, type_ids, offsets, |f| match f.name().as_str() {
4788 "long" => {
4789 Some(Arc::new(Int64Array::from(vec![1i64, -5, 42, -1, 0])) as ArrayRef)
4790 }
4791 "string" => Some(Arc::new(StringArray::from(vec!["a", "", "z"])) as ArrayRef),
4792 _ => None,
4793 });
4794 let list_offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 5, 6, 8]));
4795 expected_cols.push(Arc::new(
4796 ListArray::try_new(item_field.clone(), list_offsets, values_union, None).unwrap(),
4797 ));
4798 }
4799 {
4801 let idx = schema.index_of("map_of_union").unwrap();
4802 let dt = schema.field(idx).data_type().clone();
4803 let (entry_field, ordered) = match &dt {
4804 DataType::Map(f, ordered) => (f.clone(), *ordered),
4805 other => panic!("map_of_union must be Map, got {other:?}"),
4806 };
4807 let DataType::Struct(entry_fields) = entry_field.data_type() else {
4808 panic!("map entries must be struct")
4809 };
4810 let key_field = entry_fields[0].clone();
4811 let val_field = entry_fields[1].clone();
4812 let keys = StringArray::from(vec!["a", "b", "x", "pi"]);
4813 let rounded_pi = (std::f64::consts::PI * 100_000.0).round() / 100_000.0;
4814 let values: ArrayRef = match val_field.data_type() {
4815 DataType::Union(uf, _) => {
4816 let tid_n = tid_by_name(uf, "null");
4817 let tid_d = tid_by_name(uf, "double");
4818 let tids = vec![tid_n, tid_d, tid_d, tid_d];
4819 let offs = vec![0, 0, 1, 2];
4820 mk_dense_union(uf, tids, offs, |f| match f.name().as_str() {
4821 "null" => Some(Arc::new(NullArray::new(1)) as ArrayRef),
4822 "double" => Some(Arc::new(arrow_array::Float64Array::from(vec![
4823 2.5f64, -0.5f64, rounded_pi,
4824 ])) as ArrayRef),
4825 _ => None,
4826 })
4827 }
4828 DataType::Float64 => Arc::new(arrow_array::Float64Array::from(vec![
4829 None,
4830 Some(2.5),
4831 Some(-0.5),
4832 Some(rounded_pi),
4833 ])),
4834 other => panic!("unexpected map value type {other:?}"),
4835 };
4836 let entries = StructArray::new(
4837 Fields::from(vec![key_field.as_ref().clone(), val_field.as_ref().clone()]),
4838 vec![Arc::new(keys) as ArrayRef, values],
4839 None,
4840 );
4841 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2, 3, 3, 4]));
4842 expected_cols.push(Arc::new(MapArray::new(
4843 entry_field,
4844 offsets,
4845 entries,
4846 None,
4847 ordered,
4848 )));
4849 }
4850 {
4852 let idx = schema.index_of("record_with_union_field").unwrap();
4853 let DataType::Struct(rec_fields) = schema.field(idx).data_type() else {
4854 panic!("record_with_union_field should be Struct")
4855 };
4856 let id = Int32Array::from(vec![1, 2, 3, 4]);
4857 let u_field = rec_fields.iter().find(|f| f.name() == "u").unwrap();
4858 let DataType::Union(uf, _) = u_field.data_type() else {
4859 panic!("u must be Union")
4860 };
4861 let tid_i = tid_by_name(uf, "int");
4862 let tid_s = tid_by_name(uf, "string");
4863 let tids = vec![tid_s, tid_i, tid_i, tid_s];
4864 let offs = vec![0, 0, 1, 1];
4865 let u = mk_dense_union(uf, tids, offs, |f| match f.name().as_str() {
4866 "int" => Some(Arc::new(Int32Array::from(vec![99, 0])) as ArrayRef),
4867 "string" => Some(Arc::new(StringArray::from(vec!["one", "four"])) as ArrayRef),
4868 _ => None,
4869 });
4870 let rec = StructArray::new(rec_fields.clone(), vec![Arc::new(id) as ArrayRef, u], None);
4871 expected_cols.push(Arc::new(rec));
4872 }
4873 {
4875 let (uf, _) = get_union("union_ts_micros_utc_or_map");
4876 let tid_ts = tid_by_dt(&uf, |dt| {
4877 matches!(
4878 dt,
4879 DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, Some(_))
4880 )
4881 });
4882 let tid_map = tid_by_dt(&uf, |dt| matches!(dt, DataType::Map(_, _)));
4883 let tids = vec![tid_ts, tid_map, tid_ts, tid_map];
4884 let offs = vec![0, 0, 1, 1];
4885 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
4886 DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, tz) => {
4887 let a = TimestampMicrosecondArray::from(vec![ts_us_2024_01_01, 0i64]);
4888 Some(Arc::new(if let Some(tz) = tz {
4889 a.with_timezone(tz.clone())
4890 } else {
4891 a
4892 }) as ArrayRef)
4893 }
4894 DataType::Map(entry_field, ordered) => {
4895 let DataType::Struct(fs) = entry_field.data_type() else {
4896 panic!("map entries must be struct")
4897 };
4898 let key_field = fs[0].clone();
4899 let val_field = fs[1].clone();
4900 assert_eq!(key_field.data_type(), &DataType::Utf8);
4901 assert_eq!(val_field.data_type(), &DataType::Int64);
4902 let keys = StringArray::from(vec!["k1", "k2", "n"]);
4903 let vals = Int64Array::from(vec![1i64, 2, 0]);
4904 let entries = StructArray::new(
4905 Fields::from(vec![key_field.as_ref().clone(), val_field.as_ref().clone()]),
4906 vec![Arc::new(keys) as ArrayRef, Arc::new(vals) as ArrayRef],
4907 None,
4908 );
4909 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2, 3]));
4910 Some(Arc::new(MapArray::new(
4911 entry_field.clone(),
4912 offsets,
4913 entries,
4914 None,
4915 *ordered,
4916 )) as ArrayRef)
4917 }
4918 _ => None,
4919 });
4920 expected_cols.push(arr);
4921 }
4922 {
4924 let (uf, _) = get_union("union_ts_millis_local_or_string");
4925 let tid_ts = tid_by_dt(&uf, |dt| {
4926 matches!(
4927 dt,
4928 DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, None)
4929 )
4930 });
4931 let tid_s = tid_by_name(&uf, "string");
4932 let tids = vec![tid_s, tid_ts, tid_s, tid_s];
4933 let offs = vec![0, 0, 1, 2];
4934 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
4935 DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, None) => Some(Arc::new(
4936 TimestampMillisecondArray::from(vec![ts_ms_2024_01_01]),
4937 )
4938 as ArrayRef),
4939 DataType::Utf8 => {
4940 Some(
4941 Arc::new(StringArray::from(vec!["local midnight", "done", ""])) as ArrayRef,
4942 )
4943 }
4944 _ => None,
4945 });
4946 expected_cols.push(arr);
4947 }
4948 {
4950 let (uf, _) = get_union("union_bool_or_string");
4951 let tid_b = tid_by_name(&uf, "boolean");
4952 let tid_s = tid_by_name(&uf, "string");
4953 let tids = vec![tid_b, tid_s, tid_b, tid_s];
4954 let offs = vec![0, 0, 1, 1];
4955 let arr = mk_dense_union(&uf, tids, offs, |f| match f.name().as_str() {
4956 "boolean" => Some(Arc::new(BooleanArray::from(vec![true, false])) as ArrayRef),
4957 "string" => Some(Arc::new(StringArray::from(vec!["no", "yes"])) as ArrayRef),
4958 _ => None,
4959 });
4960 expected_cols.push(arr);
4961 }
4962 let expected = RecordBatch::try_new(schema.clone(), expected_cols).unwrap();
4963 assert_eq!(
4964 actual, expected,
4965 "full end-to-end equality for union_fields.avro"
4966 );
4967 }
4968
4969 #[test]
4970 fn test_read_zero_byte_avro_file() {
4971 let batch = read_file("test/data/zero_byte.avro", 3, false);
4972 let schema = batch.schema();
4973 assert_eq!(schema.fields().len(), 1);
4974 let field = schema.field(0);
4975 assert_eq!(field.name(), "data");
4976 assert_eq!(field.data_type(), &DataType::Binary);
4977 assert!(field.is_nullable());
4978 assert_eq!(batch.num_rows(), 3);
4979 assert_eq!(batch.num_columns(), 1);
4980 let binary_array = batch
4981 .column(0)
4982 .as_any()
4983 .downcast_ref::<BinaryArray>()
4984 .unwrap();
4985 assert!(binary_array.is_null(0));
4986 assert!(binary_array.is_valid(1));
4987 assert_eq!(binary_array.value(1), b"");
4988 assert!(binary_array.is_valid(2));
4989 assert_eq!(binary_array.value(2), b"some bytes");
4990 }
4991
4992 #[test]
4993 fn test_alltypes() {
4994 let expected = RecordBatch::try_from_iter_with_nullable([
4995 (
4996 "id",
4997 Arc::new(Int32Array::from(vec![4, 5, 6, 7, 2, 3, 0, 1])) as _,
4998 true,
4999 ),
5000 (
5001 "bool_col",
5002 Arc::new(BooleanArray::from_iter((0..8).map(|x| Some(x % 2 == 0)))) as _,
5003 true,
5004 ),
5005 (
5006 "tinyint_col",
5007 Arc::new(Int32Array::from_iter_values((0..8).map(|x| x % 2))) as _,
5008 true,
5009 ),
5010 (
5011 "smallint_col",
5012 Arc::new(Int32Array::from_iter_values((0..8).map(|x| x % 2))) as _,
5013 true,
5014 ),
5015 (
5016 "int_col",
5017 Arc::new(Int32Array::from_iter_values((0..8).map(|x| x % 2))) as _,
5018 true,
5019 ),
5020 (
5021 "bigint_col",
5022 Arc::new(Int64Array::from_iter_values((0..8).map(|x| (x % 2) * 10))) as _,
5023 true,
5024 ),
5025 (
5026 "float_col",
5027 Arc::new(Float32Array::from_iter_values(
5028 (0..8).map(|x| (x % 2) as f32 * 1.1),
5029 )) as _,
5030 true,
5031 ),
5032 (
5033 "double_col",
5034 Arc::new(Float64Array::from_iter_values(
5035 (0..8).map(|x| (x % 2) as f64 * 10.1),
5036 )) as _,
5037 true,
5038 ),
5039 (
5040 "date_string_col",
5041 Arc::new(BinaryArray::from_iter_values([
5042 [48, 51, 47, 48, 49, 47, 48, 57],
5043 [48, 51, 47, 48, 49, 47, 48, 57],
5044 [48, 52, 47, 48, 49, 47, 48, 57],
5045 [48, 52, 47, 48, 49, 47, 48, 57],
5046 [48, 50, 47, 48, 49, 47, 48, 57],
5047 [48, 50, 47, 48, 49, 47, 48, 57],
5048 [48, 49, 47, 48, 49, 47, 48, 57],
5049 [48, 49, 47, 48, 49, 47, 48, 57],
5050 ])) as _,
5051 true,
5052 ),
5053 (
5054 "string_col",
5055 Arc::new(BinaryArray::from_iter_values((0..8).map(|x| [48 + x % 2]))) as _,
5056 true,
5057 ),
5058 (
5059 "timestamp_col",
5060 Arc::new(
5061 TimestampMicrosecondArray::from_iter_values([
5062 1235865600000000, 1235865660000000, 1238544000000000, 1238544060000000, 1233446400000000, 1233446460000000, 1230768000000000, 1230768060000000, ])
5071 .with_timezone("+00:00"),
5072 ) as _,
5073 true,
5074 ),
5075 ])
5076 .unwrap();
5077
5078 for file in files() {
5079 let file = arrow_test_data(file);
5080
5081 assert_eq!(read_file(&file, 8, false), expected);
5082 assert_eq!(read_file(&file, 3, false), expected);
5083 }
5084 }
5085
5086 #[test]
5087 #[cfg(feature = "snappy")]
5089 fn test_alltypes_dictionary() {
5090 let file = "avro/alltypes_dictionary.avro";
5091 let expected = RecordBatch::try_from_iter_with_nullable([
5092 ("id", Arc::new(Int32Array::from(vec![0, 1])) as _, true),
5093 (
5094 "bool_col",
5095 Arc::new(BooleanArray::from(vec![Some(true), Some(false)])) as _,
5096 true,
5097 ),
5098 (
5099 "tinyint_col",
5100 Arc::new(Int32Array::from(vec![0, 1])) as _,
5101 true,
5102 ),
5103 (
5104 "smallint_col",
5105 Arc::new(Int32Array::from(vec![0, 1])) as _,
5106 true,
5107 ),
5108 ("int_col", Arc::new(Int32Array::from(vec![0, 1])) as _, true),
5109 (
5110 "bigint_col",
5111 Arc::new(Int64Array::from(vec![0, 10])) as _,
5112 true,
5113 ),
5114 (
5115 "float_col",
5116 Arc::new(Float32Array::from(vec![0.0, 1.1])) as _,
5117 true,
5118 ),
5119 (
5120 "double_col",
5121 Arc::new(Float64Array::from(vec![0.0, 10.1])) as _,
5122 true,
5123 ),
5124 (
5125 "date_string_col",
5126 Arc::new(BinaryArray::from_iter_values([b"01/01/09", b"01/01/09"])) as _,
5127 true,
5128 ),
5129 (
5130 "string_col",
5131 Arc::new(BinaryArray::from_iter_values([b"0", b"1"])) as _,
5132 true,
5133 ),
5134 (
5135 "timestamp_col",
5136 Arc::new(
5137 TimestampMicrosecondArray::from_iter_values([
5138 1230768000000000, 1230768060000000, ])
5141 .with_timezone("+00:00"),
5142 ) as _,
5143 true,
5144 ),
5145 ])
5146 .unwrap();
5147 let file_path = arrow_test_data(file);
5148 let batch_large = read_file(&file_path, 8, false);
5149 assert_eq!(
5150 batch_large, expected,
5151 "Decoded RecordBatch does not match for file {file}"
5152 );
5153 let batch_small = read_file(&file_path, 3, false);
5154 assert_eq!(
5155 batch_small, expected,
5156 "Decoded RecordBatch (batch size 3) does not match for file {file}"
5157 );
5158 }
5159
5160 #[test]
5161 fn test_alltypes_nulls_plain() {
5162 let file = "avro/alltypes_nulls_plain.avro";
5163 let expected = RecordBatch::try_from_iter_with_nullable([
5164 (
5165 "string_col",
5166 Arc::new(StringArray::from(vec![None::<&str>])) as _,
5167 true,
5168 ),
5169 ("int_col", Arc::new(Int32Array::from(vec![None])) as _, true),
5170 (
5171 "bool_col",
5172 Arc::new(BooleanArray::from(vec![None])) as _,
5173 true,
5174 ),
5175 (
5176 "bigint_col",
5177 Arc::new(Int64Array::from(vec![None])) as _,
5178 true,
5179 ),
5180 (
5181 "float_col",
5182 Arc::new(Float32Array::from(vec![None])) as _,
5183 true,
5184 ),
5185 (
5186 "double_col",
5187 Arc::new(Float64Array::from(vec![None])) as _,
5188 true,
5189 ),
5190 (
5191 "bytes_col",
5192 Arc::new(BinaryArray::from(vec![None::<&[u8]>])) as _,
5193 true,
5194 ),
5195 ])
5196 .unwrap();
5197 let file_path = arrow_test_data(file);
5198 let batch_large = read_file(&file_path, 8, false);
5199 assert_eq!(
5200 batch_large, expected,
5201 "Decoded RecordBatch does not match for file {file}"
5202 );
5203 let batch_small = read_file(&file_path, 3, false);
5204 assert_eq!(
5205 batch_small, expected,
5206 "Decoded RecordBatch (batch size 3) does not match for file {file}"
5207 );
5208 }
5209
5210 #[test]
5211 #[cfg(feature = "snappy")]
5213 fn test_binary() {
5214 let file = arrow_test_data("avro/binary.avro");
5215 let batch = read_file(&file, 8, false);
5216 let expected = RecordBatch::try_from_iter_with_nullable([(
5217 "foo",
5218 Arc::new(BinaryArray::from_iter_values(vec![
5219 b"\x00" as &[u8],
5220 b"\x01" as &[u8],
5221 b"\x02" as &[u8],
5222 b"\x03" as &[u8],
5223 b"\x04" as &[u8],
5224 b"\x05" as &[u8],
5225 b"\x06" as &[u8],
5226 b"\x07" as &[u8],
5227 b"\x08" as &[u8],
5228 b"\t" as &[u8],
5229 b"\n" as &[u8],
5230 b"\x0b" as &[u8],
5231 ])) as Arc<dyn Array>,
5232 true,
5233 )])
5234 .unwrap();
5235 assert_eq!(batch, expected);
5236 }
5237
5238 #[test]
5239 #[cfg(feature = "snappy")]
5241 fn test_decimal() {
5242 #[cfg(feature = "small_decimals")]
5246 let files: [(&str, DataType, HashMap<String, String>); 8] = [
5247 (
5248 "avro/fixed_length_decimal.avro",
5249 DataType::Decimal128(25, 2),
5250 HashMap::from([
5251 (
5252 "avro.namespace".to_string(),
5253 "topLevelRecord.value".to_string(),
5254 ),
5255 ("avro.name".to_string(), "fixed".to_string()),
5256 ]),
5257 ),
5258 (
5259 "avro/fixed_length_decimal_legacy.avro",
5260 DataType::Decimal64(13, 2),
5261 HashMap::from([
5262 (
5263 "avro.namespace".to_string(),
5264 "topLevelRecord.value".to_string(),
5265 ),
5266 ("avro.name".to_string(), "fixed".to_string()),
5267 ]),
5268 ),
5269 (
5270 "avro/int32_decimal.avro",
5271 DataType::Decimal32(4, 2),
5272 HashMap::from([
5273 (
5274 "avro.namespace".to_string(),
5275 "topLevelRecord.value".to_string(),
5276 ),
5277 ("avro.name".to_string(), "fixed".to_string()),
5278 ]),
5279 ),
5280 (
5281 "avro/int64_decimal.avro",
5282 DataType::Decimal64(10, 2),
5283 HashMap::from([
5284 (
5285 "avro.namespace".to_string(),
5286 "topLevelRecord.value".to_string(),
5287 ),
5288 ("avro.name".to_string(), "fixed".to_string()),
5289 ]),
5290 ),
5291 (
5292 "test/data/int256_decimal.avro",
5293 DataType::Decimal256(76, 10),
5294 HashMap::new(),
5295 ),
5296 (
5297 "test/data/fixed256_decimal.avro",
5298 DataType::Decimal256(76, 10),
5299 HashMap::from([("avro.name".to_string(), "Decimal256Fixed".to_string())]),
5300 ),
5301 (
5302 "test/data/fixed_length_decimal_legacy_32.avro",
5303 DataType::Decimal32(9, 2),
5304 HashMap::from([("avro.name".to_string(), "Decimal32FixedLegacy".to_string())]),
5305 ),
5306 (
5307 "test/data/int128_decimal.avro",
5308 DataType::Decimal128(38, 2),
5309 HashMap::new(),
5310 ),
5311 ];
5312 #[cfg(not(feature = "small_decimals"))]
5313 let files: [(&str, DataType, HashMap<String, String>); 8] = [
5314 (
5315 "avro/fixed_length_decimal.avro",
5316 DataType::Decimal128(25, 2),
5317 HashMap::from([
5318 (
5319 "avro.namespace".to_string(),
5320 "topLevelRecord.value".to_string(),
5321 ),
5322 ("avro.name".to_string(), "fixed".to_string()),
5323 ]),
5324 ),
5325 (
5326 "avro/fixed_length_decimal_legacy.avro",
5327 DataType::Decimal128(13, 2),
5328 HashMap::from([
5329 (
5330 "avro.namespace".to_string(),
5331 "topLevelRecord.value".to_string(),
5332 ),
5333 ("avro.name".to_string(), "fixed".to_string()),
5334 ]),
5335 ),
5336 (
5337 "avro/int32_decimal.avro",
5338 DataType::Decimal128(4, 2),
5339 HashMap::from([
5340 (
5341 "avro.namespace".to_string(),
5342 "topLevelRecord.value".to_string(),
5343 ),
5344 ("avro.name".to_string(), "fixed".to_string()),
5345 ]),
5346 ),
5347 (
5348 "avro/int64_decimal.avro",
5349 DataType::Decimal128(10, 2),
5350 HashMap::from([
5351 (
5352 "avro.namespace".to_string(),
5353 "topLevelRecord.value".to_string(),
5354 ),
5355 ("avro.name".to_string(), "fixed".to_string()),
5356 ]),
5357 ),
5358 (
5359 "test/data/int256_decimal.avro",
5360 DataType::Decimal256(76, 10),
5361 HashMap::new(),
5362 ),
5363 (
5364 "test/data/fixed256_decimal.avro",
5365 DataType::Decimal256(76, 10),
5366 HashMap::from([("avro.name".to_string(), "Decimal256Fixed".to_string())]),
5367 ),
5368 (
5369 "test/data/fixed_length_decimal_legacy_32.avro",
5370 DataType::Decimal128(9, 2),
5371 HashMap::from([("avro.name".to_string(), "Decimal32FixedLegacy".to_string())]),
5372 ),
5373 (
5374 "test/data/int128_decimal.avro",
5375 DataType::Decimal128(38, 2),
5376 HashMap::new(),
5377 ),
5378 ];
5379 for (file, expected_dt, mut metadata) in files {
5380 let (DataType::Decimal32(precision, scale)
5381 | DataType::Decimal64(precision, scale)
5382 | DataType::Decimal128(precision, scale)
5383 | DataType::Decimal256(precision, scale)) = expected_dt
5384 else {
5385 unreachable!("Unexpected decimal type in test inputs")
5386 };
5387 assert!(scale >= 0, "test data uses non-negative scales only");
5388 let scale_u32 = scale as u32;
5389 let file_path: String = if file.starts_with("avro/") {
5390 arrow_test_data(file)
5391 } else {
5392 std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR"))
5393 .join(file)
5394 .to_string_lossy()
5395 .into_owned()
5396 };
5397 let pow10 = 10i128.pow(scale_u32);
5398 let values_i128: Vec<i128> = (1..=24).map(|n| (n as i128) * pow10).collect();
5399 let build_expected = |dt: &DataType, values: &[i128]| -> ArrayRef {
5400 match *dt {
5401 #[cfg(feature = "small_decimals")]
5402 DataType::Decimal32(p, s) => {
5403 let it = values.iter().map(|&v| v as i32);
5404 Arc::new(
5405 Decimal32Array::from_iter_values(it)
5406 .with_precision_and_scale(p, s)
5407 .unwrap(),
5408 )
5409 }
5410 #[cfg(feature = "small_decimals")]
5411 DataType::Decimal64(p, s) => {
5412 let it = values.iter().map(|&v| v as i64);
5413 Arc::new(
5414 Decimal64Array::from_iter_values(it)
5415 .with_precision_and_scale(p, s)
5416 .unwrap(),
5417 )
5418 }
5419 DataType::Decimal128(p, s) => {
5420 let it = values.iter().copied();
5421 Arc::new(
5422 Decimal128Array::from_iter_values(it)
5423 .with_precision_and_scale(p, s)
5424 .unwrap(),
5425 )
5426 }
5427 DataType::Decimal256(p, s) => {
5428 let it = values.iter().map(|&v| i256::from_i128(v));
5429 Arc::new(
5430 Decimal256Array::from_iter_values(it)
5431 .with_precision_and_scale(p, s)
5432 .unwrap(),
5433 )
5434 }
5435 _ => unreachable!("Unexpected decimal type in test"),
5436 }
5437 };
5438 let actual_batch = read_file(&file_path, 8, false);
5439 let actual_nullable = actual_batch.schema().field(0).is_nullable();
5440 let expected_array = build_expected(&expected_dt, &values_i128);
5441 metadata.insert("precision".to_string(), precision.to_string());
5442 metadata.insert("scale".to_string(), scale.to_string());
5443 let field =
5444 Field::new("value", expected_dt.clone(), actual_nullable).with_metadata(metadata);
5445 let expected_schema = Arc::new(Schema::new(vec![field]));
5446 let expected_batch =
5447 RecordBatch::try_new(expected_schema.clone(), vec![expected_array]).unwrap();
5448 assert_eq!(
5449 actual_batch, expected_batch,
5450 "Decoded RecordBatch does not match for {file}"
5451 );
5452 let actual_batch_small = read_file(&file_path, 3, false);
5453 assert_eq!(
5454 actual_batch_small, expected_batch,
5455 "Decoded RecordBatch does not match for {file} with batch size 3"
5456 );
5457 }
5458 }
5459
5460 #[test]
5461 fn test_read_duration_logical_types_feature_toggle() -> Result<(), ArrowError> {
5462 let file_path = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR"))
5463 .join("test/data/duration_logical_types.avro")
5464 .to_string_lossy()
5465 .into_owned();
5466
5467 let actual_batch = read_file(&file_path, 4, false);
5468
5469 let expected_batch = {
5470 #[cfg(feature = "avro_custom_types")]
5471 {
5472 let schema = Arc::new(Schema::new(vec![
5473 Field::new(
5474 "duration_time_nanos",
5475 DataType::Duration(TimeUnit::Nanosecond),
5476 false,
5477 ),
5478 Field::new(
5479 "duration_time_micros",
5480 DataType::Duration(TimeUnit::Microsecond),
5481 false,
5482 ),
5483 Field::new(
5484 "duration_time_millis",
5485 DataType::Duration(TimeUnit::Millisecond),
5486 false,
5487 ),
5488 Field::new(
5489 "duration_time_seconds",
5490 DataType::Duration(TimeUnit::Second),
5491 false,
5492 ),
5493 ]));
5494
5495 let nanos = Arc::new(PrimitiveArray::<DurationNanosecondType>::from(vec![
5496 10, 20, 30, 40,
5497 ])) as ArrayRef;
5498 let micros = Arc::new(PrimitiveArray::<DurationMicrosecondType>::from(vec![
5499 100, 200, 300, 400,
5500 ])) as ArrayRef;
5501 let millis = Arc::new(PrimitiveArray::<DurationMillisecondType>::from(vec![
5502 1000, 2000, 3000, 4000,
5503 ])) as ArrayRef;
5504 let seconds = Arc::new(PrimitiveArray::<DurationSecondType>::from(vec![1, 2, 3, 4]))
5505 as ArrayRef;
5506
5507 RecordBatch::try_new(schema, vec![nanos, micros, millis, seconds])?
5508 }
5509 #[cfg(not(feature = "avro_custom_types"))]
5510 {
5511 let schema = Arc::new(Schema::new(vec![
5512 Field::new("duration_time_nanos", DataType::Int64, false)
5513 .with_metadata([("logicalType", "arrow.duration-nanos")]),
5514 Field::new("duration_time_micros", DataType::Int64, false)
5515 .with_metadata([("logicalType", "arrow.duration-micros")]),
5516 Field::new("duration_time_millis", DataType::Int64, false)
5517 .with_metadata([("logicalType", "arrow.duration-millis")]),
5518 Field::new("duration_time_seconds", DataType::Int64, false)
5519 .with_metadata([("logicalType", "arrow.duration-seconds")]),
5520 ]));
5521
5522 let nanos =
5523 Arc::new(PrimitiveArray::<Int64Type>::from(vec![10, 20, 30, 40])) as ArrayRef;
5524 let micros = Arc::new(PrimitiveArray::<Int64Type>::from(vec![100, 200, 300, 400]))
5525 as ArrayRef;
5526 let millis = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
5527 1000, 2000, 3000, 4000,
5528 ])) as ArrayRef;
5529 let seconds =
5530 Arc::new(PrimitiveArray::<Int64Type>::from(vec![1, 2, 3, 4])) as ArrayRef;
5531
5532 RecordBatch::try_new(schema, vec![nanos, micros, millis, seconds])?
5533 }
5534 };
5535
5536 assert_eq!(actual_batch, expected_batch);
5537
5538 Ok(())
5539 }
5540
5541 #[test]
5542 #[cfg(feature = "snappy")]
5544 fn test_dict_pages_offset_zero() {
5545 let file = arrow_test_data("avro/dict-page-offset-zero.avro");
5546 let batch = read_file(&file, 32, false);
5547 let num_rows = batch.num_rows();
5548 let expected_field = Int32Array::from(vec![Some(1552); num_rows]);
5549 let expected = RecordBatch::try_from_iter_with_nullable([(
5550 "l_partkey",
5551 Arc::new(expected_field) as Arc<dyn Array>,
5552 true,
5553 )])
5554 .unwrap();
5555 assert_eq!(batch, expected);
5556 }
5557
5558 #[test]
5559 #[cfg(feature = "snappy")]
5561 fn test_list_columns() {
5562 let file = arrow_test_data("avro/list_columns.avro");
5563 let mut int64_list_builder = ListBuilder::new(Int64Builder::new());
5564 {
5565 {
5566 let values = int64_list_builder.values();
5567 values.append_value(1);
5568 values.append_value(2);
5569 values.append_value(3);
5570 }
5571 int64_list_builder.append(true);
5572 }
5573 {
5574 {
5575 let values = int64_list_builder.values();
5576 values.append_null();
5577 values.append_value(1);
5578 }
5579 int64_list_builder.append(true);
5580 }
5581 {
5582 {
5583 let values = int64_list_builder.values();
5584 values.append_value(4);
5585 }
5586 int64_list_builder.append(true);
5587 }
5588 let int64_list = int64_list_builder.finish();
5589 let mut utf8_list_builder = ListBuilder::new(StringBuilder::new());
5590 {
5591 {
5592 let values = utf8_list_builder.values();
5593 values.append_value("abc");
5594 values.append_value("efg");
5595 values.append_value("hij");
5596 }
5597 utf8_list_builder.append(true);
5598 }
5599 {
5600 utf8_list_builder.append(false);
5601 }
5602 {
5603 {
5604 let values = utf8_list_builder.values();
5605 values.append_value("efg");
5606 values.append_null();
5607 values.append_value("hij");
5608 values.append_value("xyz");
5609 }
5610 utf8_list_builder.append(true);
5611 }
5612 let utf8_list = utf8_list_builder.finish();
5613 let expected = RecordBatch::try_from_iter_with_nullable([
5614 ("int64_list", Arc::new(int64_list) as Arc<dyn Array>, true),
5615 ("utf8_list", Arc::new(utf8_list) as Arc<dyn Array>, true),
5616 ])
5617 .unwrap();
5618 let batch = read_file(&file, 8, false);
5619 assert_eq!(batch, expected);
5620 }
5621
5622 #[test]
5623 #[cfg(feature = "snappy")]
5624 fn test_nested_lists() {
5625 use arrow_data::ArrayDataBuilder;
5626 let file = arrow_test_data("avro/nested_lists.snappy.avro");
5627 let inner_values = StringArray::from(vec![
5628 Some("a"),
5629 Some("b"),
5630 Some("c"),
5631 Some("d"),
5632 Some("a"),
5633 Some("b"),
5634 Some("c"),
5635 Some("d"),
5636 Some("e"),
5637 Some("a"),
5638 Some("b"),
5639 Some("c"),
5640 Some("d"),
5641 Some("e"),
5642 Some("f"),
5643 ]);
5644 let inner_offsets = Buffer::from_slice_ref([0, 2, 3, 3, 4, 6, 8, 8, 9, 11, 13, 14, 14, 15]);
5645 let inner_validity = [
5646 true, true, false, true, true, true, false, true, true, true, true, false, true,
5647 ];
5648 let inner_null_buffer = Buffer::from_iter(inner_validity.iter().copied());
5649 let inner_field = Field::new("item", DataType::Utf8, true);
5650 let inner_list_data = ArrayDataBuilder::new(DataType::List(Arc::new(inner_field)))
5651 .len(13)
5652 .add_buffer(inner_offsets)
5653 .add_child_data(inner_values.to_data())
5654 .null_bit_buffer(Some(inner_null_buffer))
5655 .build()
5656 .unwrap();
5657 let inner_list_array = ListArray::from(inner_list_data);
5658 let middle_offsets = Buffer::from_slice_ref([0, 2, 4, 6, 8, 11, 13]);
5659 let middle_validity = [true; 6];
5660 let middle_null_buffer = Buffer::from_iter(middle_validity.iter().copied());
5661 let middle_field = Field::new("item", inner_list_array.data_type().clone(), true);
5662 let middle_list_data = ArrayDataBuilder::new(DataType::List(Arc::new(middle_field)))
5663 .len(6)
5664 .add_buffer(middle_offsets)
5665 .add_child_data(inner_list_array.to_data())
5666 .null_bit_buffer(Some(middle_null_buffer))
5667 .build()
5668 .unwrap();
5669 let middle_list_array = ListArray::from(middle_list_data);
5670 let outer_offsets = Buffer::from_slice_ref([0, 2, 4, 6]);
5671 let outer_null_buffer = Buffer::from_slice_ref([0b111]); let outer_field = Field::new("item", middle_list_array.data_type().clone(), true);
5673 let outer_list_data = ArrayDataBuilder::new(DataType::List(Arc::new(outer_field)))
5674 .len(3)
5675 .add_buffer(outer_offsets)
5676 .add_child_data(middle_list_array.to_data())
5677 .null_bit_buffer(Some(outer_null_buffer))
5678 .build()
5679 .unwrap();
5680 let a_expected = ListArray::from(outer_list_data);
5681 let b_expected = Int32Array::from(vec![1, 1, 1]);
5682 let expected = RecordBatch::try_from_iter_with_nullable([
5683 ("a", Arc::new(a_expected) as Arc<dyn Array>, true),
5684 ("b", Arc::new(b_expected) as Arc<dyn Array>, true),
5685 ])
5686 .unwrap();
5687 let left = read_file(&file, 8, false);
5688 assert_eq!(left, expected, "Mismatch for batch size=8");
5689 let left_small = read_file(&file, 3, false);
5690 assert_eq!(left_small, expected, "Mismatch for batch size=3");
5691 }
5692
5693 #[test]
5694 fn test_simple() {
5695 let tests = [
5696 ("avro/simple_enum.avro", 4, build_expected_enum(), 2),
5697 ("avro/simple_fixed.avro", 2, build_expected_fixed(), 1),
5698 ];
5699
5700 fn build_expected_enum() -> RecordBatch {
5701 let keys_f1 = Int32Array::from(vec![0, 1, 2, 3]);
5703 let vals_f1 = StringArray::from(vec!["a", "b", "c", "d"]);
5704 let f1_dict =
5705 DictionaryArray::<Int32Type>::try_new(keys_f1, Arc::new(vals_f1)).unwrap();
5706 let keys_f2 = Int32Array::from(vec![2, 3, 0, 1]);
5707 let vals_f2 = StringArray::from(vec!["e", "f", "g", "h"]);
5708 let f2_dict =
5709 DictionaryArray::<Int32Type>::try_new(keys_f2, Arc::new(vals_f2)).unwrap();
5710 let keys_f3 = Int32Array::from(vec![Some(1), Some(2), None, Some(0)]);
5711 let vals_f3 = StringArray::from(vec!["i", "j", "k"]);
5712 let f3_dict =
5713 DictionaryArray::<Int32Type>::try_new(keys_f3, Arc::new(vals_f3)).unwrap();
5714 let dict_type =
5715 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
5716 let mut md_f1 = HashMap::new();
5717 md_f1.insert(
5718 AVRO_ENUM_SYMBOLS_METADATA_KEY.to_string(),
5719 r#"["a","b","c","d"]"#.to_string(),
5720 );
5721 md_f1.insert(AVRO_NAME_METADATA_KEY.to_string(), "enum1".to_string());
5722 md_f1.insert(AVRO_NAMESPACE_METADATA_KEY.to_string(), "ns1".to_string());
5723 let f1_field = Field::new("f1", dict_type.clone(), false).with_metadata(md_f1);
5724 let mut md_f2 = HashMap::new();
5725 md_f2.insert(
5726 AVRO_ENUM_SYMBOLS_METADATA_KEY.to_string(),
5727 r#"["e","f","g","h"]"#.to_string(),
5728 );
5729 md_f2.insert(AVRO_NAME_METADATA_KEY.to_string(), "enum2".to_string());
5730 md_f2.insert(AVRO_NAMESPACE_METADATA_KEY.to_string(), "ns2".to_string());
5731 let f2_field = Field::new("f2", dict_type.clone(), false).with_metadata(md_f2);
5732 let mut md_f3 = HashMap::new();
5733 md_f3.insert(
5734 AVRO_ENUM_SYMBOLS_METADATA_KEY.to_string(),
5735 r#"["i","j","k"]"#.to_string(),
5736 );
5737 md_f3.insert(AVRO_NAME_METADATA_KEY.to_string(), "enum3".to_string());
5738 md_f3.insert(AVRO_NAMESPACE_METADATA_KEY.to_string(), "ns1".to_string());
5739 let f3_field = Field::new("f3", dict_type.clone(), true).with_metadata(md_f3);
5740 let expected_schema = Arc::new(Schema::new(vec![f1_field, f2_field, f3_field]));
5741 RecordBatch::try_new(
5742 expected_schema,
5743 vec![
5744 Arc::new(f1_dict) as Arc<dyn Array>,
5745 Arc::new(f2_dict) as Arc<dyn Array>,
5746 Arc::new(f3_dict) as Arc<dyn Array>,
5747 ],
5748 )
5749 .unwrap()
5750 }
5751
5752 fn build_expected_fixed() -> RecordBatch {
5753 let f1 =
5754 FixedSizeBinaryArray::try_from_iter(vec![b"abcde", b"12345"].into_iter()).unwrap();
5755 let f2 =
5756 FixedSizeBinaryArray::try_from_iter(vec![b"fghijklmno", b"1234567890"].into_iter())
5757 .unwrap();
5758 let f3 = FixedSizeBinaryArray::try_from_sparse_iter_with_size(
5759 vec![Some(b"ABCDEF" as &[u8]), None].into_iter(),
5760 6,
5761 )
5762 .unwrap();
5763
5764 let mut md_f1 = HashMap::new();
5766 md_f1.insert(
5767 crate::schema::AVRO_NAME_METADATA_KEY.to_string(),
5768 "fixed1".to_string(),
5769 );
5770 md_f1.insert(
5771 crate::schema::AVRO_NAMESPACE_METADATA_KEY.to_string(),
5772 "ns1".to_string(),
5773 );
5774
5775 let mut md_f2 = HashMap::new();
5776 md_f2.insert(
5777 crate::schema::AVRO_NAME_METADATA_KEY.to_string(),
5778 "fixed2".to_string(),
5779 );
5780 md_f2.insert(
5781 crate::schema::AVRO_NAMESPACE_METADATA_KEY.to_string(),
5782 "ns2".to_string(),
5783 );
5784
5785 let mut md_f3 = HashMap::new();
5786 md_f3.insert(
5787 crate::schema::AVRO_NAME_METADATA_KEY.to_string(),
5788 "fixed3".to_string(),
5789 );
5790 md_f3.insert(
5791 crate::schema::AVRO_NAMESPACE_METADATA_KEY.to_string(),
5792 "ns1".to_string(),
5793 );
5794
5795 let expected_schema = Arc::new(Schema::new(vec![
5796 Field::new("f1", DataType::FixedSizeBinary(5), false).with_metadata(md_f1),
5797 Field::new("f2", DataType::FixedSizeBinary(10), false).with_metadata(md_f2),
5798 Field::new("f3", DataType::FixedSizeBinary(6), true).with_metadata(md_f3),
5799 ]));
5800
5801 RecordBatch::try_new(
5802 expected_schema,
5803 vec![
5804 Arc::new(f1) as Arc<dyn Array>,
5805 Arc::new(f2) as Arc<dyn Array>,
5806 Arc::new(f3) as Arc<dyn Array>,
5807 ],
5808 )
5809 .unwrap()
5810 }
5811 for (file_name, batch_size, expected, alt_batch_size) in tests {
5812 let file = arrow_test_data(file_name);
5813 let actual = read_file(&file, batch_size, false);
5814 assert_eq!(actual, expected);
5815 let actual2 = read_file(&file, alt_batch_size, false);
5816 assert_eq!(actual2, expected);
5817 }
5818 }
5819
5820 #[test]
5821 #[cfg(feature = "snappy")]
5822 fn test_single_nan() {
5823 let file = arrow_test_data("avro/single_nan.avro");
5824 let actual = read_file(&file, 1, false);
5825 use arrow_array::Float64Array;
5826 let schema = Arc::new(Schema::new(vec![Field::new(
5827 "mycol",
5828 DataType::Float64,
5829 true,
5830 )]));
5831 let col = Float64Array::from(vec![None]);
5832 let expected = RecordBatch::try_new(schema, vec![Arc::new(col)]).unwrap();
5833 assert_eq!(actual, expected);
5834 let actual2 = read_file(&file, 2, false);
5835 assert_eq!(actual2, expected);
5836 }
5837
5838 #[test]
5839 fn test_duration_uuid() {
5840 let batch = read_file("test/data/duration_uuid.avro", 4, false);
5841 let schema = batch.schema();
5842 let fields = schema.fields();
5843 assert_eq!(fields.len(), 2);
5844 assert_eq!(fields[0].name(), "duration_field");
5845 assert_eq!(
5846 fields[0].data_type(),
5847 &DataType::Interval(IntervalUnit::MonthDayNano)
5848 );
5849 assert_eq!(fields[1].name(), "uuid_field");
5850 assert_eq!(fields[1].data_type(), &DataType::FixedSizeBinary(16));
5851 assert_eq!(batch.num_rows(), 4);
5852 assert_eq!(batch.num_columns(), 2);
5853 let duration_array = batch
5854 .column(0)
5855 .as_any()
5856 .downcast_ref::<IntervalMonthDayNanoArray>()
5857 .unwrap();
5858 let expected_duration_array: IntervalMonthDayNanoArray = [
5859 Some(IntervalMonthDayNanoType::make_value(1, 15, 500_000_000)),
5860 Some(IntervalMonthDayNanoType::make_value(0, 5, 2_500_000_000)),
5861 Some(IntervalMonthDayNanoType::make_value(2, 0, 0)),
5862 Some(IntervalMonthDayNanoType::make_value(12, 31, 999_000_000)),
5863 ]
5864 .iter()
5865 .copied()
5866 .collect();
5867 assert_eq!(&expected_duration_array, duration_array);
5868 let uuid_array = batch
5869 .column(1)
5870 .as_any()
5871 .downcast_ref::<FixedSizeBinaryArray>()
5872 .unwrap();
5873 let expected_uuid_array = FixedSizeBinaryArray::try_from_sparse_iter_with_size(
5874 [
5875 Some([
5876 0xfe, 0x7b, 0xc3, 0x0b, 0x4c, 0xe8, 0x4c, 0x5e, 0xb6, 0x7c, 0x22, 0x34, 0xa2,
5877 0xd3, 0x8e, 0x66,
5878 ]),
5879 Some([
5880 0xb3, 0x3f, 0x2a, 0xd7, 0x97, 0xb4, 0x4d, 0xe1, 0x8b, 0xfe, 0x94, 0x94, 0x1d,
5881 0x60, 0x15, 0x6e,
5882 ]),
5883 Some([
5884 0x5f, 0x74, 0x92, 0x64, 0x07, 0x4b, 0x40, 0x05, 0x84, 0xbf, 0x11, 0x5e, 0xa8,
5885 0x4e, 0xd2, 0x0a,
5886 ]),
5887 Some([
5888 0x08, 0x26, 0xcc, 0x06, 0xd2, 0xe3, 0x45, 0x99, 0xb4, 0xad, 0xaf, 0x5f, 0xa6,
5889 0x90, 0x5c, 0xdb,
5890 ]),
5891 ]
5892 .into_iter(),
5893 16,
5894 )
5895 .unwrap();
5896 assert_eq!(&expected_uuid_array, uuid_array);
5897 }
5898
5899 #[test]
5900 #[cfg(feature = "snappy")]
5901 fn test_datapage_v2() {
5902 let file = arrow_test_data("avro/datapage_v2.snappy.avro");
5903 let batch = read_file(&file, 8, false);
5904 let a = StringArray::from(vec![
5905 Some("abc"),
5906 Some("abc"),
5907 Some("abc"),
5908 None,
5909 Some("abc"),
5910 ]);
5911 let b = Int32Array::from(vec![Some(1), Some(2), Some(3), Some(4), Some(5)]);
5912 let c = Float64Array::from(vec![Some(2.0), Some(3.0), Some(4.0), Some(5.0), Some(2.0)]);
5913 let d = BooleanArray::from(vec![
5914 Some(true),
5915 Some(true),
5916 Some(true),
5917 Some(false),
5918 Some(true),
5919 ]);
5920 let e_values = Int32Array::from(vec![
5921 Some(1),
5922 Some(2),
5923 Some(3),
5924 Some(1),
5925 Some(2),
5926 Some(3),
5927 Some(1),
5928 Some(2),
5929 ]);
5930 let e_offsets = OffsetBuffer::new(ScalarBuffer::from(vec![0i32, 3, 3, 3, 6, 8]));
5931 let e_validity = Some(NullBuffer::from(vec![true, false, false, true, true]));
5932 let field_e = Arc::new(Field::new("item", DataType::Int32, true));
5933 let e = ListArray::new(field_e, e_offsets, Arc::new(e_values), e_validity);
5934 let expected = RecordBatch::try_from_iter_with_nullable([
5935 ("a", Arc::new(a) as Arc<dyn Array>, true),
5936 ("b", Arc::new(b) as Arc<dyn Array>, true),
5937 ("c", Arc::new(c) as Arc<dyn Array>, true),
5938 ("d", Arc::new(d) as Arc<dyn Array>, true),
5939 ("e", Arc::new(e) as Arc<dyn Array>, true),
5940 ])
5941 .unwrap();
5942 assert_eq!(batch, expected);
5943 }
5944
5945 #[test]
5946 fn test_nested_records() {
5947 let f1_f1_1 = StringArray::from(vec!["aaa", "bbb"]);
5948 let f1_f1_2 = Int32Array::from(vec![10, 20]);
5949 let rounded_pi = (std::f64::consts::PI * 100.0).round() / 100.0;
5950 let f1_f1_3_1 = Float64Array::from(vec![rounded_pi, rounded_pi]);
5951 let f1_f1_3 = StructArray::from(vec![(
5952 Arc::new(Field::new("f1_3_1", DataType::Float64, false)),
5953 Arc::new(f1_f1_3_1) as Arc<dyn Array>,
5954 )]);
5955 let mut f1_3_md: HashMap<String, String> = HashMap::new();
5957 f1_3_md.insert(AVRO_NAMESPACE_METADATA_KEY.to_string(), "ns3".to_string());
5958 f1_3_md.insert(AVRO_NAME_METADATA_KEY.to_string(), "record3".to_string());
5959 let f1_expected = StructArray::from(vec![
5960 (
5961 Arc::new(Field::new("f1_1", DataType::Utf8, false)),
5962 Arc::new(f1_f1_1) as Arc<dyn Array>,
5963 ),
5964 (
5965 Arc::new(Field::new("f1_2", DataType::Int32, false)),
5966 Arc::new(f1_f1_2) as Arc<dyn Array>,
5967 ),
5968 (
5969 Arc::new(
5970 Field::new(
5971 "f1_3",
5972 DataType::Struct(Fields::from(vec![Field::new(
5973 "f1_3_1",
5974 DataType::Float64,
5975 false,
5976 )])),
5977 false,
5978 )
5979 .with_metadata(f1_3_md),
5980 ),
5981 Arc::new(f1_f1_3) as Arc<dyn Array>,
5982 ),
5983 ]);
5984 let f2_fields = [
5985 Field::new("f2_1", DataType::Boolean, false),
5986 Field::new("f2_2", DataType::Float32, false),
5987 ];
5988 let f2_struct_builder = StructBuilder::new(
5989 f2_fields
5990 .iter()
5991 .map(|f| Arc::new(f.clone()))
5992 .collect::<Vec<Arc<Field>>>(),
5993 vec![
5994 Box::new(BooleanBuilder::new()) as Box<dyn arrow_array::builder::ArrayBuilder>,
5995 Box::new(Float32Builder::new()) as Box<dyn arrow_array::builder::ArrayBuilder>,
5996 ],
5997 );
5998 let mut f2_list_builder = ListBuilder::new(f2_struct_builder);
5999 {
6000 let struct_builder = f2_list_builder.values();
6001 struct_builder.append(true);
6002 {
6003 let b = struct_builder.field_builder::<BooleanBuilder>(0).unwrap();
6004 b.append_value(true);
6005 }
6006 {
6007 let b = struct_builder.field_builder::<Float32Builder>(1).unwrap();
6008 b.append_value(1.2_f32);
6009 }
6010 struct_builder.append(true);
6011 {
6012 let b = struct_builder.field_builder::<BooleanBuilder>(0).unwrap();
6013 b.append_value(true);
6014 }
6015 {
6016 let b = struct_builder.field_builder::<Float32Builder>(1).unwrap();
6017 b.append_value(2.2_f32);
6018 }
6019 f2_list_builder.append(true);
6020 }
6021 {
6022 let struct_builder = f2_list_builder.values();
6023 struct_builder.append(true);
6024 {
6025 let b = struct_builder.field_builder::<BooleanBuilder>(0).unwrap();
6026 b.append_value(false);
6027 }
6028 {
6029 let b = struct_builder.field_builder::<Float32Builder>(1).unwrap();
6030 b.append_value(10.2_f32);
6031 }
6032 f2_list_builder.append(true);
6033 }
6034
6035 let list_array_with_nullable_items = f2_list_builder.finish();
6036 let mut f2_item_md: HashMap<String, String> = HashMap::new();
6038 f2_item_md.insert(AVRO_NAME_METADATA_KEY.to_string(), "record4".to_string());
6039 f2_item_md.insert(AVRO_NAMESPACE_METADATA_KEY.to_string(), "ns4".to_string());
6040 let item_field = Arc::new(
6041 Field::new(
6042 "item",
6043 list_array_with_nullable_items.values().data_type().clone(),
6044 false, )
6046 .with_metadata(f2_item_md),
6047 );
6048 let list_data_type = DataType::List(item_field);
6049 let f2_array_data = list_array_with_nullable_items
6050 .to_data()
6051 .into_builder()
6052 .data_type(list_data_type)
6053 .build()
6054 .unwrap();
6055 let f2_expected = ListArray::from(f2_array_data);
6056 let mut f3_struct_builder = StructBuilder::new(
6057 vec![Arc::new(Field::new("f3_1", DataType::Utf8, false))],
6058 vec![Box::new(StringBuilder::new()) as Box<dyn ArrayBuilder>],
6059 );
6060 f3_struct_builder.append(true);
6061 {
6062 let b = f3_struct_builder.field_builder::<StringBuilder>(0).unwrap();
6063 b.append_value("xyz");
6064 }
6065 f3_struct_builder.append(false);
6066 {
6067 let b = f3_struct_builder.field_builder::<StringBuilder>(0).unwrap();
6068 b.append_null();
6069 }
6070 let f3_expected = f3_struct_builder.finish();
6071 let f4_fields = [Field::new("f4_1", DataType::Int64, false)];
6072 let f4_struct_builder = StructBuilder::new(
6073 f4_fields
6074 .iter()
6075 .map(|f| Arc::new(f.clone()))
6076 .collect::<Vec<Arc<Field>>>(),
6077 vec![Box::new(Int64Builder::new()) as Box<dyn arrow_array::builder::ArrayBuilder>],
6078 );
6079 let mut f4_list_builder = ListBuilder::new(f4_struct_builder);
6080 {
6081 let struct_builder = f4_list_builder.values();
6082 struct_builder.append(true);
6083 {
6084 let b = struct_builder.field_builder::<Int64Builder>(0).unwrap();
6085 b.append_value(200);
6086 }
6087 struct_builder.append(false);
6088 {
6089 let b = struct_builder.field_builder::<Int64Builder>(0).unwrap();
6090 b.append_null();
6091 }
6092 f4_list_builder.append(true);
6093 }
6094 {
6095 let struct_builder = f4_list_builder.values();
6096 struct_builder.append(false);
6097 {
6098 let b = struct_builder.field_builder::<Int64Builder>(0).unwrap();
6099 b.append_null();
6100 }
6101 struct_builder.append(true);
6102 {
6103 let b = struct_builder.field_builder::<Int64Builder>(0).unwrap();
6104 b.append_value(300);
6105 }
6106 f4_list_builder.append(true);
6107 }
6108 let f4_expected = f4_list_builder.finish();
6109 let mut f4_item_md: HashMap<String, String> = HashMap::new();
6111 f4_item_md.insert(AVRO_NAMESPACE_METADATA_KEY.to_string(), "ns6".to_string());
6112 f4_item_md.insert(AVRO_NAME_METADATA_KEY.to_string(), "record6".to_string());
6113 let f4_item_field = Arc::new(
6114 Field::new("item", f4_expected.values().data_type().clone(), true)
6115 .with_metadata(f4_item_md),
6116 );
6117 let f4_list_data_type = DataType::List(f4_item_field);
6118 let f4_array_data = f4_expected
6119 .to_data()
6120 .into_builder()
6121 .data_type(f4_list_data_type)
6122 .build()
6123 .unwrap();
6124 let f4_expected = ListArray::from(f4_array_data);
6125 let mut f1_md: HashMap<String, String> = HashMap::new();
6127 f1_md.insert(AVRO_NAME_METADATA_KEY.to_string(), "record2".to_string());
6128 f1_md.insert(AVRO_NAMESPACE_METADATA_KEY.to_string(), "ns2".to_string());
6129 let mut f3_md: HashMap<String, String> = HashMap::new();
6130 f3_md.insert(AVRO_NAMESPACE_METADATA_KEY.to_string(), "ns5".to_string());
6131 f3_md.insert(AVRO_NAME_METADATA_KEY.to_string(), "record5".to_string());
6132 let expected_schema = Schema::new(vec![
6133 Field::new("f1", f1_expected.data_type().clone(), false).with_metadata(f1_md),
6134 Field::new("f2", f2_expected.data_type().clone(), false),
6135 Field::new("f3", f3_expected.data_type().clone(), true).with_metadata(f3_md),
6136 Field::new("f4", f4_expected.data_type().clone(), false),
6137 ]);
6138 let expected = RecordBatch::try_new(
6139 Arc::new(expected_schema),
6140 vec![
6141 Arc::new(f1_expected) as Arc<dyn Array>,
6142 Arc::new(f2_expected) as Arc<dyn Array>,
6143 Arc::new(f3_expected) as Arc<dyn Array>,
6144 Arc::new(f4_expected) as Arc<dyn Array>,
6145 ],
6146 )
6147 .unwrap();
6148 let file = arrow_test_data("avro/nested_records.avro");
6149 let batch_large = read_file(&file, 8, false);
6150 assert_eq!(
6151 batch_large, expected,
6152 "Decoded RecordBatch does not match expected data for nested records (batch size 8)"
6153 );
6154 let batch_small = read_file(&file, 3, false);
6155 assert_eq!(
6156 batch_small, expected,
6157 "Decoded RecordBatch does not match expected data for nested records (batch size 3)"
6158 );
6159 }
6160
6161 #[test]
6162 #[cfg(feature = "snappy")]
6164 fn test_repeated_no_annotation() {
6165 use arrow_data::ArrayDataBuilder;
6166 let file = arrow_test_data("avro/repeated_no_annotation.avro");
6167 let batch_large = read_file(&file, 8, false);
6168 let id_array = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
6170 let number_array = Int64Array::from(vec![
6172 Some(5555555555),
6173 Some(1111111111),
6174 Some(1111111111),
6175 Some(2222222222),
6176 Some(3333333333),
6177 ]);
6178 let kind_array =
6179 StringArray::from(vec![None, Some("home"), Some("home"), None, Some("mobile")]);
6180 let phone_fields = Fields::from(vec![
6181 Field::new("number", DataType::Int64, true),
6182 Field::new("kind", DataType::Utf8, true),
6183 ]);
6184 let phone_struct_data = ArrayDataBuilder::new(DataType::Struct(phone_fields))
6185 .len(5)
6186 .child_data(vec![number_array.into_data(), kind_array.into_data()])
6187 .build()
6188 .unwrap();
6189 let phone_struct_array = StructArray::from(phone_struct_data);
6190 let phone_list_offsets = Buffer::from_slice_ref([0i32, 0, 0, 0, 1, 2, 5]);
6192 let phone_list_validity = Buffer::from_iter([false, false, true, true, true, true]);
6193 let mut phone_item_md = HashMap::new();
6195 phone_item_md.insert(AVRO_NAME_METADATA_KEY.to_string(), "phone".to_string());
6196 phone_item_md.insert(
6197 AVRO_NAMESPACE_METADATA_KEY.to_string(),
6198 "topLevelRecord.phoneNumbers".to_string(),
6199 );
6200 let phone_item_field = Field::new("item", phone_struct_array.data_type().clone(), true)
6201 .with_metadata(phone_item_md);
6202 let phone_list_data = ArrayDataBuilder::new(DataType::List(Arc::new(phone_item_field)))
6203 .len(6)
6204 .add_buffer(phone_list_offsets)
6205 .null_bit_buffer(Some(phone_list_validity))
6206 .child_data(vec![phone_struct_array.into_data()])
6207 .build()
6208 .unwrap();
6209 let phone_list_array = ListArray::from(phone_list_data);
6210 let phone_numbers_validity = Buffer::from_iter([false, false, true, true, true, true]);
6212 let phone_numbers_field = Field::new("phone", phone_list_array.data_type().clone(), true);
6213 let phone_numbers_struct_data =
6214 ArrayDataBuilder::new(DataType::Struct(Fields::from(vec![phone_numbers_field])))
6215 .len(6)
6216 .null_bit_buffer(Some(phone_numbers_validity))
6217 .child_data(vec![phone_list_array.into_data()])
6218 .build()
6219 .unwrap();
6220 let phone_numbers_struct_array = StructArray::from(phone_numbers_struct_data);
6221 let mut phone_numbers_md = HashMap::new();
6223 phone_numbers_md.insert(
6224 AVRO_NAME_METADATA_KEY.to_string(),
6225 "phoneNumbers".to_string(),
6226 );
6227 phone_numbers_md.insert(
6228 AVRO_NAMESPACE_METADATA_KEY.to_string(),
6229 "topLevelRecord".to_string(),
6230 );
6231 let id_field = Field::new("id", DataType::Int32, true);
6232 let phone_numbers_schema_field = Field::new(
6233 "phoneNumbers",
6234 phone_numbers_struct_array.data_type().clone(),
6235 true,
6236 )
6237 .with_metadata(phone_numbers_md);
6238 let expected_schema = Schema::new(vec![id_field, phone_numbers_schema_field]);
6239 let expected = RecordBatch::try_new(
6241 Arc::new(expected_schema),
6242 vec![
6243 Arc::new(id_array) as _,
6244 Arc::new(phone_numbers_struct_array) as _,
6245 ],
6246 )
6247 .unwrap();
6248 assert_eq!(batch_large, expected, "Mismatch for batch_size=8");
6249 let batch_small = read_file(&file, 3, false);
6250 assert_eq!(batch_small, expected, "Mismatch for batch_size=3");
6251 }
6252
6253 #[test]
6254 #[cfg(feature = "snappy")]
6256 fn test_nonnullable_impala() {
6257 let file = arrow_test_data("avro/nonnullable.impala.avro");
6258 let id = Int64Array::from(vec![Some(8)]);
6259 let mut int_array_builder = ListBuilder::new(Int32Builder::new());
6260 {
6261 let vb = int_array_builder.values();
6262 vb.append_value(-1);
6263 }
6264 int_array_builder.append(true); let int_array = int_array_builder.finish();
6266 let mut iaa_builder = ListBuilder::new(ListBuilder::new(Int32Builder::new()));
6267 {
6268 let inner_list_builder = iaa_builder.values();
6269 {
6270 let vb = inner_list_builder.values();
6271 vb.append_value(-1);
6272 vb.append_value(-2);
6273 }
6274 inner_list_builder.append(true);
6275 inner_list_builder.append(true);
6276 }
6277 iaa_builder.append(true);
6278 let int_array_array = iaa_builder.finish();
6279 let field_names = MapFieldNames {
6280 entry: Field::MAP_ENTRIES_FIELD_DEFAULT_NAME.to_string(),
6281 key: Field::MAP_KEY_FIELD_DEFAULT_NAME.to_string(),
6282 value: Field::MAP_VALUE_FIELD_DEFAULT_NAME.to_string(),
6283 };
6284 let mut int_map_builder =
6285 MapBuilder::new(Some(field_names), StringBuilder::new(), Int32Builder::new());
6286 {
6287 let (keys, vals) = int_map_builder.entries();
6288 keys.append_value("k1");
6289 vals.append_value(-1);
6290 }
6291 int_map_builder.append(true).unwrap(); let int_map = int_map_builder.finish();
6293 let field_names2 = MapFieldNames {
6294 entry: Field::MAP_ENTRIES_FIELD_DEFAULT_NAME.to_string(),
6295 key: Field::MAP_KEY_FIELD_DEFAULT_NAME.to_string(),
6296 value: Field::MAP_VALUE_FIELD_DEFAULT_NAME.to_string(),
6297 };
6298 let mut ima_builder = ListBuilder::new(MapBuilder::new(
6299 Some(field_names2),
6300 StringBuilder::new(),
6301 Int32Builder::new(),
6302 ));
6303 {
6304 let map_builder = ima_builder.values();
6305 map_builder.append(true).unwrap();
6306 {
6307 let (keys, vals) = map_builder.entries();
6308 keys.append_value("k1");
6309 vals.append_value(1);
6310 }
6311 map_builder.append(true).unwrap();
6312 map_builder.append(true).unwrap();
6313 map_builder.append(true).unwrap();
6314 }
6315 ima_builder.append(true);
6316 let int_map_array_ = ima_builder.finish();
6317 let meta_nested_struct: HashMap<String, String> = [
6319 ("avro.name", "nested_Struct"),
6320 ("avro.namespace", "topLevelRecord"),
6321 ]
6322 .into_iter()
6323 .map(|(k, v)| (k.to_string(), v.to_string()))
6324 .collect();
6325 let meta_c: HashMap<String, String> = [
6326 ("avro.name", "c"),
6327 ("avro.namespace", "topLevelRecord.nested_Struct"),
6328 ]
6329 .into_iter()
6330 .map(|(k, v)| (k.to_string(), v.to_string()))
6331 .collect();
6332 let meta_d_item_struct: HashMap<String, String> = [
6333 ("avro.name", "D"),
6334 ("avro.namespace", "topLevelRecord.nested_Struct.c"),
6335 ]
6336 .into_iter()
6337 .map(|(k, v)| (k.to_string(), v.to_string()))
6338 .collect();
6339 let meta_g_value: HashMap<String, String> = [
6340 ("avro.name", "G"),
6341 ("avro.namespace", "topLevelRecord.nested_Struct"),
6342 ]
6343 .into_iter()
6344 .map(|(k, v)| (k.to_string(), v.to_string()))
6345 .collect();
6346 let meta_h: HashMap<String, String> = [
6347 ("avro.name", "h"),
6348 ("avro.namespace", "topLevelRecord.nested_Struct.G"),
6349 ]
6350 .into_iter()
6351 .map(|(k, v)| (k.to_string(), v.to_string()))
6352 .collect();
6353 let ef_struct_field = Arc::new(
6355 Field::new(
6356 "item",
6357 DataType::Struct(
6358 vec![
6359 Field::new("e", DataType::Int32, true),
6360 Field::new("f", DataType::Utf8, true),
6361 ]
6362 .into(),
6363 ),
6364 true,
6365 )
6366 .with_metadata(meta_d_item_struct.clone()),
6367 );
6368 let d_inner_list_field = Arc::new(Field::new(
6369 "item",
6370 DataType::List(ef_struct_field.clone()),
6371 true,
6372 ));
6373 let d_field = Field::new("D", DataType::List(d_inner_list_field.clone()), true);
6374 let i_list_field = Arc::new(Field::new("item", DataType::Float64, true));
6376 let i_field = Field::new("i", DataType::List(i_list_field.clone()), true);
6377 let h_field = Field::new("h", DataType::Struct(vec![i_field.clone()].into()), true)
6379 .with_metadata(meta_h.clone());
6380 let g_value_struct_field = Field::new(
6382 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6383 DataType::Struct(vec![h_field.clone()].into()),
6384 true,
6385 )
6386 .with_metadata(meta_g_value.clone());
6387 let entries_struct_field = Field::new(
6389 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
6390 DataType::Struct(
6391 vec![
6392 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
6393 g_value_struct_field.clone(),
6394 ]
6395 .into(),
6396 ),
6397 false,
6398 );
6399 let a_field = Arc::new(Field::new("a", DataType::Int32, true));
6401 let b_field = Arc::new(Field::new(
6402 "B",
6403 DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
6404 true,
6405 ));
6406 let c_field = Arc::new(
6407 Field::new("c", DataType::Struct(vec![d_field.clone()].into()), true)
6408 .with_metadata(meta_c.clone()),
6409 );
6410 let g_field = Arc::new(Field::new(
6411 "G",
6412 DataType::Map(Arc::new(entries_struct_field.clone()), false),
6413 true,
6414 ));
6415 let mut nested_sb = StructBuilder::new(
6417 vec![
6418 a_field.clone(),
6419 b_field.clone(),
6420 c_field.clone(),
6421 g_field.clone(),
6422 ],
6423 vec![
6424 Box::new(Int32Builder::new()),
6425 Box::new(ListBuilder::new(Int32Builder::new())),
6426 {
6427 Box::new(StructBuilder::new(
6429 vec![Arc::new(d_field.clone())],
6430 vec![Box::new({
6431 let ef_struct_builder = StructBuilder::new(
6432 vec![
6433 Arc::new(Field::new("e", DataType::Int32, true)),
6434 Arc::new(Field::new("f", DataType::Utf8, true)),
6435 ],
6436 vec![
6437 Box::new(Int32Builder::new()),
6438 Box::new(StringBuilder::new()),
6439 ],
6440 );
6441 let list_of_ef = ListBuilder::new(ef_struct_builder)
6443 .with_field(ef_struct_field.clone());
6444 ListBuilder::new(list_of_ef)
6446 })],
6447 ))
6448 },
6449 {
6450 let map_field_names = MapFieldNames {
6451 entry: Field::MAP_ENTRIES_FIELD_DEFAULT_NAME.to_string(),
6452 key: Field::MAP_KEY_FIELD_DEFAULT_NAME.to_string(),
6453 value: Field::MAP_VALUE_FIELD_DEFAULT_NAME.to_string(),
6454 };
6455 let i_list_builder = ListBuilder::new(Float64Builder::new());
6456 let h_struct_builder = StructBuilder::new(
6457 vec![Arc::new(Field::new(
6458 "i",
6459 DataType::List(i_list_field.clone()),
6460 true,
6461 ))],
6462 vec![Box::new(i_list_builder)],
6463 );
6464 let g_value_builder = StructBuilder::new(
6465 vec![Arc::new(
6466 Field::new("h", DataType::Struct(vec![i_field.clone()].into()), true)
6467 .with_metadata(meta_h.clone()),
6468 )],
6469 vec![Box::new(h_struct_builder)],
6470 );
6471 let map_builder = MapBuilder::new(
6473 Some(map_field_names),
6474 StringBuilder::new(),
6475 g_value_builder,
6476 )
6477 .with_values_field(Arc::new(
6478 Field::new(
6479 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
6480 DataType::Struct(vec![h_field.clone()].into()),
6481 true,
6482 )
6483 .with_metadata(meta_g_value.clone()),
6484 ));
6485
6486 Box::new(map_builder)
6487 },
6488 ],
6489 );
6490 nested_sb.append(true);
6491 {
6492 let a_builder = nested_sb.field_builder::<Int32Builder>(0).unwrap();
6493 a_builder.append_value(-1);
6494 }
6495 {
6496 let b_builder = nested_sb
6497 .field_builder::<ListBuilder<Int32Builder>>(1)
6498 .unwrap();
6499 {
6500 let vb = b_builder.values();
6501 vb.append_value(-1);
6502 }
6503 b_builder.append(true);
6504 }
6505 {
6506 let c_struct_builder = nested_sb.field_builder::<StructBuilder>(2).unwrap();
6507 c_struct_builder.append(true);
6508 let d_list_builder = c_struct_builder
6509 .field_builder::<ListBuilder<ListBuilder<StructBuilder>>>(0)
6510 .unwrap();
6511 {
6512 let sub_list_builder = d_list_builder.values();
6513 {
6514 let ef_struct = sub_list_builder.values();
6515 ef_struct.append(true);
6516 {
6517 let e_b = ef_struct.field_builder::<Int32Builder>(0).unwrap();
6518 e_b.append_value(-1);
6519 let f_b = ef_struct.field_builder::<StringBuilder>(1).unwrap();
6520 f_b.append_value("nonnullable");
6521 }
6522 sub_list_builder.append(true);
6523 }
6524 d_list_builder.append(true);
6525 }
6526 }
6527 {
6528 let g_map_builder = nested_sb
6529 .field_builder::<MapBuilder<StringBuilder, StructBuilder>>(3)
6530 .unwrap();
6531 g_map_builder.append(true).unwrap();
6532 }
6533 let nested_struct = nested_sb.finish();
6534 let schema = Arc::new(arrow_schema::Schema::new(vec![
6535 Field::new("ID", id.data_type().clone(), true),
6536 Field::new("Int_Array", int_array.data_type().clone(), true),
6537 Field::new("int_array_array", int_array_array.data_type().clone(), true),
6538 Field::new("Int_Map", int_map.data_type().clone(), true),
6539 Field::new("int_map_array", int_map_array_.data_type().clone(), true),
6540 Field::new("nested_Struct", nested_struct.data_type().clone(), true)
6541 .with_metadata(meta_nested_struct.clone()),
6542 ]));
6543 let expected = RecordBatch::try_new(
6544 schema,
6545 vec![
6546 Arc::new(id) as Arc<dyn Array>,
6547 Arc::new(int_array),
6548 Arc::new(int_array_array),
6549 Arc::new(int_map),
6550 Arc::new(int_map_array_),
6551 Arc::new(nested_struct),
6552 ],
6553 )
6554 .unwrap();
6555 let batch_large = read_file(&file, 8, false);
6556 assert_eq!(batch_large, expected, "Mismatch for batch_size=8");
6557 let batch_small = read_file(&file, 3, false);
6558 assert_eq!(batch_small, expected, "Mismatch for batch_size=3");
6559 }
6560
6561 #[test]
6562 fn test_nonnullable_impala_strict() {
6563 let file = arrow_test_data("avro/nonnullable.impala.avro");
6564 let err = read_file_strict(&file, 8, false).unwrap_err();
6565 assert!(err.to_string().contains(
6566 "Found Avro union of the form ['T','null'], which is disallowed in strict_mode"
6567 ));
6568 }
6569
6570 #[test]
6571 #[cfg(feature = "snappy")]
6573 fn test_nullable_impala() {
6574 let file = arrow_test_data("avro/nullable.impala.avro");
6575 let batch1 = read_file(&file, 3, false);
6576 let batch2 = read_file(&file, 8, false);
6577 assert_eq!(batch1, batch2);
6578 let batch = batch1;
6579 assert_eq!(batch.num_rows(), 7);
6580 let id_array = batch
6581 .column(0)
6582 .as_any()
6583 .downcast_ref::<Int64Array>()
6584 .expect("id column should be an Int64Array");
6585 let expected_ids = [1, 2, 3, 4, 5, 6, 7];
6586 for (i, &expected_id) in expected_ids.iter().enumerate() {
6587 assert_eq!(id_array.value(i), expected_id, "Mismatch in id at row {i}");
6588 }
6589 let int_array = batch
6590 .column(1)
6591 .as_any()
6592 .downcast_ref::<ListArray>()
6593 .expect("int_array column should be a ListArray");
6594 {
6595 let offsets = int_array.value_offsets();
6596 let start = offsets[0] as usize;
6597 let end = offsets[1] as usize;
6598 let values = int_array
6599 .values()
6600 .as_any()
6601 .downcast_ref::<Int32Array>()
6602 .expect("Values of int_array should be an Int32Array");
6603 let row0: Vec<Option<i32>> = (start..end).map(|i| Some(values.value(i))).collect();
6604 assert_eq!(
6605 row0,
6606 vec![Some(1), Some(2), Some(3)],
6607 "Mismatch in int_array row 0"
6608 );
6609 }
6610 let nested_struct = batch
6611 .column(5)
6612 .as_any()
6613 .downcast_ref::<StructArray>()
6614 .expect("nested_struct column should be a StructArray");
6615 let a_array = nested_struct
6616 .column_by_name("A")
6617 .expect("Field A should exist in nested_struct")
6618 .as_any()
6619 .downcast_ref::<Int32Array>()
6620 .expect("Field A should be an Int32Array");
6621 assert_eq!(a_array.value(0), 1, "Mismatch in nested_struct.A at row 0");
6622 assert!(
6623 !a_array.is_valid(1),
6624 "Expected null in nested_struct.A at row 1"
6625 );
6626 assert!(
6627 !a_array.is_valid(3),
6628 "Expected null in nested_struct.A at row 3"
6629 );
6630 assert_eq!(a_array.value(6), 7, "Mismatch in nested_struct.A at row 6");
6631 }
6632
6633 #[test]
6634 fn test_nullable_impala_strict() {
6635 let file = arrow_test_data("avro/nullable.impala.avro");
6636 let err = read_file_strict(&file, 8, false).unwrap_err();
6637 assert!(err.to_string().contains(
6638 "Found Avro union of the form ['T','null'], which is disallowed in strict_mode"
6639 ));
6640 }
6641
6642 #[test]
6643 fn test_nested_record_type_reuse() {
6644 let batch = read_file("test/data/nested_record_reuse.avro", 8, false);
6670 let schema = batch.schema();
6671
6672 assert_eq!(schema.fields().len(), 3);
6674 let fields = schema.fields();
6675 assert_eq!(fields[0].name(), "nested");
6676 assert_eq!(fields[1].name(), "nestedRecord");
6677 assert_eq!(fields[2].name(), "nestedArray");
6678 assert!(matches!(fields[0].data_type(), DataType::Struct(_)));
6679 assert!(matches!(fields[1].data_type(), DataType::Struct(_)));
6680 assert!(matches!(fields[2].data_type(), DataType::List(_)));
6681
6682 if let DataType::Struct(nested_fields) = fields[0].data_type() {
6684 assert_eq!(nested_fields.len(), 1);
6685 assert_eq!(nested_fields[0].name(), "nested_int");
6686 assert_eq!(nested_fields[0].data_type(), &DataType::Int32);
6687 }
6688
6689 assert_eq!(fields[0].data_type(), fields[1].data_type());
6691 if let DataType::List(array_field) = fields[2].data_type() {
6692 assert_eq!(array_field.data_type(), fields[0].data_type());
6693 }
6694
6695 assert_eq!(batch.num_rows(), 2);
6697 assert_eq!(batch.num_columns(), 3);
6698
6699 let nested_col = batch
6701 .column(0)
6702 .as_any()
6703 .downcast_ref::<StructArray>()
6704 .unwrap();
6705 let nested_int_array = nested_col
6706 .column_by_name("nested_int")
6707 .unwrap()
6708 .as_any()
6709 .downcast_ref::<Int32Array>()
6710 .unwrap();
6711 assert_eq!(nested_int_array.value(0), 42);
6712 assert_eq!(nested_int_array.value(1), 99);
6713
6714 let nested_record_col = batch
6716 .column(1)
6717 .as_any()
6718 .downcast_ref::<StructArray>()
6719 .unwrap();
6720 let nested_record_int_array = nested_record_col
6721 .column_by_name("nested_int")
6722 .unwrap()
6723 .as_any()
6724 .downcast_ref::<Int32Array>()
6725 .unwrap();
6726 assert_eq!(nested_record_int_array.value(0), 100);
6727 assert_eq!(nested_record_int_array.value(1), 200);
6728
6729 let nested_array_col = batch
6731 .column(2)
6732 .as_any()
6733 .downcast_ref::<ListArray>()
6734 .unwrap();
6735 assert_eq!(nested_array_col.len(), 2);
6736 let first_array_struct = nested_array_col.value(0);
6737 let first_array_struct_array = first_array_struct
6738 .as_any()
6739 .downcast_ref::<StructArray>()
6740 .unwrap();
6741 let first_array_int_values = first_array_struct_array
6742 .column_by_name("nested_int")
6743 .unwrap()
6744 .as_any()
6745 .downcast_ref::<Int32Array>()
6746 .unwrap();
6747 assert_eq!(first_array_int_values.len(), 3);
6748 assert_eq!(first_array_int_values.value(0), 1);
6749 assert_eq!(first_array_int_values.value(1), 2);
6750 assert_eq!(first_array_int_values.value(2), 3);
6751 }
6752
6753 #[test]
6754 fn test_enum_type_reuse() {
6755 let batch = read_file("test/data/enum_reuse.avro", 8, false);
6778 let schema = batch.schema();
6779
6780 assert_eq!(schema.fields().len(), 3);
6782 let fields = schema.fields();
6783 assert_eq!(fields[0].name(), "status");
6784 assert_eq!(fields[1].name(), "backupStatus");
6785 assert_eq!(fields[2].name(), "statusHistory");
6786 assert!(matches!(fields[0].data_type(), DataType::Dictionary(_, _)));
6787 assert!(matches!(fields[1].data_type(), DataType::Dictionary(_, _)));
6788 assert!(matches!(fields[2].data_type(), DataType::List(_)));
6789
6790 if let DataType::Dictionary(key_type, value_type) = fields[0].data_type() {
6791 assert_eq!(key_type.as_ref(), &DataType::Int32);
6792 assert_eq!(value_type.as_ref(), &DataType::Utf8);
6793 }
6794
6795 assert_eq!(fields[0].data_type(), fields[1].data_type());
6797 if let DataType::List(array_field) = fields[2].data_type() {
6798 assert_eq!(array_field.data_type(), fields[0].data_type());
6799 }
6800
6801 assert_eq!(batch.num_rows(), 2);
6803 assert_eq!(batch.num_columns(), 3);
6804
6805 let status_col = batch
6807 .column(0)
6808 .as_any()
6809 .downcast_ref::<DictionaryArray<Int32Type>>()
6810 .unwrap();
6811 let status_values = status_col
6812 .values()
6813 .as_any()
6814 .downcast_ref::<StringArray>()
6815 .unwrap();
6816
6817 assert_eq!(status_values.value(status_col.key(0).unwrap()), "ACTIVE");
6819 assert_eq!(status_values.value(status_col.key(1).unwrap()), "PENDING");
6820
6821 let backup_status_col = batch
6823 .column(1)
6824 .as_any()
6825 .downcast_ref::<DictionaryArray<Int32Type>>()
6826 .unwrap();
6827 let backup_status_values = backup_status_col
6828 .values()
6829 .as_any()
6830 .downcast_ref::<StringArray>()
6831 .unwrap();
6832
6833 assert_eq!(
6835 backup_status_values.value(backup_status_col.key(0).unwrap()),
6836 "INACTIVE"
6837 );
6838 assert_eq!(
6839 backup_status_values.value(backup_status_col.key(1).unwrap()),
6840 "ACTIVE"
6841 );
6842
6843 let status_history_col = batch
6845 .column(2)
6846 .as_any()
6847 .downcast_ref::<ListArray>()
6848 .unwrap();
6849 assert_eq!(status_history_col.len(), 2);
6850
6851 let first_array_dict = status_history_col.value(0);
6853 let first_array_dict_array = first_array_dict
6854 .as_any()
6855 .downcast_ref::<DictionaryArray<Int32Type>>()
6856 .unwrap();
6857 let first_array_values = first_array_dict_array
6858 .values()
6859 .as_any()
6860 .downcast_ref::<StringArray>()
6861 .unwrap();
6862
6863 assert_eq!(first_array_dict_array.len(), 3);
6865 assert_eq!(
6866 first_array_values.value(first_array_dict_array.key(0).unwrap()),
6867 "PENDING"
6868 );
6869 assert_eq!(
6870 first_array_values.value(first_array_dict_array.key(1).unwrap()),
6871 "ACTIVE"
6872 );
6873 assert_eq!(
6874 first_array_values.value(first_array_dict_array.key(2).unwrap()),
6875 "INACTIVE"
6876 );
6877 }
6878
6879 #[test]
6880 fn test_bad_varint_bug_nullable_array_items() {
6881 use flate2::read::GzDecoder;
6882 use std::io::Read;
6883 let manifest_dir = env!("CARGO_MANIFEST_DIR");
6884 let gz_path = format!("{manifest_dir}/test/data/bad-varint-bug.avro.gz");
6885 let gz_file = File::open(&gz_path).expect("test file should exist");
6886 let mut decoder = GzDecoder::new(gz_file);
6887 let mut avro_bytes = Vec::new();
6888 decoder
6889 .read_to_end(&mut avro_bytes)
6890 .expect("should decompress");
6891 let reader_arrow_schema = Schema::new(vec![Field::new(
6892 "int_array",
6893 DataType::List(Arc::new(Field::new("element", DataType::Int32, true))),
6894 true,
6895 )])
6896 .with_metadata(HashMap::from([("avro.name".into(), "table".into())]));
6897 let reader_schema = AvroSchema::try_from(&reader_arrow_schema)
6898 .expect("should convert Arrow schema to Avro");
6899 let mut reader = ReaderBuilder::new()
6900 .with_reader_schema(reader_schema)
6901 .build(Cursor::new(avro_bytes))
6902 .expect("should build reader");
6903 let batch = reader
6904 .next()
6905 .expect("should have one batch")
6906 .expect("reading should succeed without bad varint error");
6907 assert_eq!(batch.num_rows(), 1);
6908 let list_col = batch
6909 .column(0)
6910 .as_any()
6911 .downcast_ref::<ListArray>()
6912 .expect("should be ListArray");
6913 assert_eq!(list_col.len(), 1);
6914 let values = list_col.values();
6915 let int_values = values.as_primitive::<Int32Type>();
6916 assert_eq!(int_values.len(), 2);
6917 assert_eq!(int_values.value(0), 1);
6918 assert_eq!(int_values.value(1), 2);
6919 }
6920
6921 #[test]
6922 fn test_nested_record_field_addition() {
6923 let file = arrow_test_data("avro/nested_records.avro");
6924
6925 let reader_schema = AvroSchema::new(
6937 r#"
6938 {
6939 "type": "record",
6940 "name": "record1",
6941 "namespace": "ns1",
6942 "fields": [
6943 {
6944 "name": "f1",
6945 "type": [
6946 "null",
6947 {
6948 "type": "record",
6949 "name": "record2",
6950 "namespace": "ns2",
6951 "fields": [
6952 {
6953 "name": "f1_1",
6954 "type": "string"
6955 },
6956 {
6957 "name": "f1_2",
6958 "type": "int"
6959 },
6960 {
6961 "name": "f1_3",
6962 "type": {
6963 "type": "record",
6964 "name": "record3",
6965 "namespace": "ns3",
6966 "fields": [
6967 {
6968 "name": "f1_3_1",
6969 "type": "double"
6970 }
6971 ]
6972 }
6973 },
6974 {
6975 "name": "f1_4",
6976 "type": ["null", "int"],
6977 "default": null
6978 }
6979 ]
6980 }
6981 ]
6982 },
6983 {
6984 "name": "f2",
6985 "type": {
6986 "type": "array",
6987 "items": {
6988 "type": "record",
6989 "name": "record4",
6990 "namespace": "ns4",
6991 "fields": [
6992 {
6993 "name": "f2_1",
6994 "type": "boolean"
6995 },
6996 {
6997 "name": "f2_2",
6998 "type": "float"
6999 },
7000 {
7001 "name": "f2_3",
7002 "type": ["null", "int"],
7003 "default": 42
7004 }
7005 ]
7006 }
7007 }
7008 },
7009 {
7010 "name": "f3",
7011 "type": [
7012 "null",
7013 {
7014 "type": "record",
7015 "name": "record5",
7016 "namespace": "ns5",
7017 "fields": [
7018 {
7019 "name": "f3_0",
7020 "type": "string",
7021 "default": "lorem ipsum"
7022 },
7023 {
7024 "name": "f3_1",
7025 "type": "string"
7026 }
7027 ]
7028 }
7029 ],
7030 "default": null
7031 },
7032 {
7033 "name": "f4",
7034 "type": {
7035 "type": "array",
7036 "items": [
7037 "null",
7038 {
7039 "type": "record",
7040 "name": "record6",
7041 "namespace": "ns6",
7042 "fields": [
7043 {
7044 "name": "f4_1",
7045 "type": "long"
7046 }
7047 ]
7048 }
7049 ]
7050 }
7051 }
7052 ]
7053 }
7054 "#
7055 .to_string(),
7056 );
7057
7058 let file = File::open(&file).unwrap();
7059 let mut reader = ReaderBuilder::new()
7060 .with_reader_schema(reader_schema)
7061 .build(BufReader::new(file))
7062 .expect("reader with evolved reader schema should be built successfully");
7063
7064 let batch = reader
7065 .next()
7066 .expect("should have at least one batch")
7067 .expect("reading should succeed");
7068
7069 assert!(batch.num_rows() > 0);
7070
7071 let schema = batch.schema();
7072
7073 let f1_field = schema.field_with_name("f1").expect("f1 field should exist");
7074 if let DataType::Struct(f1_fields) = f1_field.data_type() {
7075 let (_, f1_4) = f1_fields
7076 .find("f1_4")
7077 .expect("f1_4 field should be present in record2");
7078 assert!(f1_4.is_nullable(), "f1_4 should be nullable");
7079 assert_eq!(f1_4.data_type(), &DataType::Int32, "f1_4 should be Int32");
7080 assert_eq!(
7081 f1_4.metadata().get("avro.field.default"),
7082 Some(&"null".to_string()),
7083 "f1_4 should have null default value in metadata"
7084 );
7085 } else {
7086 panic!("f1 should be a struct");
7087 }
7088
7089 let f2_field = schema.field_with_name("f2").expect("f2 field should exist");
7090 if let DataType::List(f2_items_field) = f2_field.data_type() {
7091 if let DataType::Struct(f2_items_fields) = f2_items_field.data_type() {
7092 let (_, f2_3) = f2_items_fields
7093 .find("f2_3")
7094 .expect("f2_3 field should be present in record4");
7095 assert!(f2_3.is_nullable(), "f2_3 should be nullable");
7096 assert_eq!(f2_3.data_type(), &DataType::Int32, "f2_3 should be Int32");
7097 assert_eq!(
7098 f2_3.metadata().get("avro.field.default"),
7099 Some(&"42".to_string()),
7100 "f2_3 should have 42 default value in metadata"
7101 );
7102 } else {
7103 panic!("f2 array items should be a struct");
7104 }
7105 } else {
7106 panic!("f2 should be a list");
7107 }
7108
7109 let f3_field = schema.field_with_name("f3").expect("f3 field should exist");
7110 assert!(f3_field.is_nullable(), "f3 should be nullable");
7111 if let DataType::Struct(f3_fields) = f3_field.data_type() {
7112 let (_, f3_0) = f3_fields
7113 .find("f3_0")
7114 .expect("f3_0 field should be present in record5");
7115 assert!(!f3_0.is_nullable(), "f3_0 should be non-nullable");
7116 assert_eq!(f3_0.data_type(), &DataType::Utf8, "f3_0 should be a string");
7117 assert_eq!(
7118 f3_0.metadata().get("avro.field.default"),
7119 Some(&"\"lorem ipsum\"".to_string()),
7120 "f3_0 should have \"lorem ipsum\" default value in metadata"
7121 );
7122 } else {
7123 panic!("f3 should be a struct");
7124 }
7125
7126 let num_rows = batch.num_rows();
7128
7129 let f1_array = batch
7131 .column_by_name("f1")
7132 .expect("f1 column should exist")
7133 .as_struct();
7134 let f1_4_array = f1_array
7135 .column_by_name("f1_4")
7136 .expect("f1_4 column should exist in f1 struct")
7137 .as_primitive::<Int32Type>();
7138
7139 assert_eq!(f1_4_array.null_count(), num_rows);
7140
7141 let f2_array = batch
7142 .column_by_name("f2")
7143 .expect("f2 column should exist")
7144 .as_list::<i32>();
7145
7146 for i in 0..num_rows {
7147 assert!(!f2_array.is_null(i));
7148 let f2_value = f2_array.value(i);
7149 let f2_record_array = f2_value.as_struct();
7150 let f2_3_array = f2_record_array
7151 .column_by_name("f2_3")
7152 .expect("f2_3 column should exist in f2 array items")
7153 .as_primitive::<Int32Type>();
7154
7155 for j in 0..f2_3_array.len() {
7156 assert!(!f2_3_array.is_null(j));
7157 assert_eq!(f2_3_array.value(j), 42);
7158 }
7159 }
7160
7161 let f3_array = batch
7162 .column_by_name("f3")
7163 .expect("f3 column should exist")
7164 .as_struct();
7165 let f3_0_array = f3_array
7166 .column_by_name("f3_0")
7167 .expect("f3_0 column should exist in f3 struct")
7168 .as_string::<i32>();
7169
7170 for i in 0..num_rows {
7171 if !f3_array.is_null(i) {
7173 assert!(!f3_0_array.is_null(i));
7174 assert_eq!(f3_0_array.value(i), "lorem ipsum");
7175 }
7176 }
7177 }
7178
7179 fn corrupt_first_block_payload_byte(
7180 mut bytes: Vec<u8>,
7181 field_offset: usize,
7182 expected_original: u8,
7183 replacement: u8,
7184 ) -> Vec<u8> {
7185 let mut header_decoder = HeaderDecoder::default();
7186 let header_len = header_decoder.decode(&bytes).expect("decode header");
7187 assert!(header_decoder.flush().is_some(), "decode complete header");
7188
7189 let mut cursor = &bytes[header_len..];
7190 let (_, count_len) = crate::reader::vlq::read_varint(cursor).expect("decode block count");
7191 cursor = &cursor[count_len..];
7192 let (_, size_len) = crate::reader::vlq::read_varint(cursor).expect("decode block size");
7193 let data_start = header_len + count_len + size_len;
7194 let target = data_start + field_offset;
7195
7196 assert!(
7197 target < bytes.len(),
7198 "target byte offset {target} out of bounds for input length {}",
7199 bytes.len()
7200 );
7201 assert_eq!(
7202 bytes[target], expected_original,
7203 "unexpected original byte at payload offset {field_offset}"
7204 );
7205 bytes[target] = replacement;
7206 bytes
7207 }
7208
7209 #[test]
7210 fn ocf_projection_rejects_overflowing_varint_in_skipped_long_field() {
7211 let writer_schema = Schema::new(vec![
7215 Field::new("bad_long", DataType::Int64, false),
7216 Field::new("keep", DataType::Int32, false),
7217 ]);
7218 let batch = RecordBatch::try_new(
7219 Arc::new(writer_schema.clone()),
7220 vec![
7221 Arc::new(Int64Array::from(vec![i64::MIN])) as ArrayRef,
7222 Arc::new(Int32Array::from(vec![7])) as ArrayRef,
7223 ],
7224 )
7225 .expect("build writer batch");
7226 let bytes = write_ocf(&writer_schema, &[batch]);
7227 let mutated = corrupt_first_block_payload_byte(bytes, 9, 0x01, 0x02);
7228
7229 let err = ReaderBuilder::new()
7230 .build(Cursor::new(mutated.clone()))
7231 .expect("build full reader")
7232 .collect::<Result<Vec<_>, _>>()
7233 .expect_err("full decode should reject malformed varint");
7234 assert!(matches!(err, ArrowError::AvroError(_)));
7235 assert!(err.to_string().contains("bad varint"));
7236
7237 let err = ReaderBuilder::new()
7238 .with_projection(vec![1])
7239 .build(Cursor::new(mutated))
7240 .expect("build projected reader")
7241 .collect::<Result<Vec<_>, _>>()
7242 .expect_err("projection must also reject malformed skipped varint");
7243 assert!(matches!(err, ArrowError::AvroError(_)));
7244 assert!(err.to_string().contains("bad varint"));
7245 }
7246
7247 #[test]
7248 fn ocf_projection_rejects_i32_overflow_in_skipped_int_field() {
7249 let writer_schema = Schema::new(vec![
7253 Field::new("bad_int", DataType::Int32, false),
7254 Field::new("keep", DataType::Int64, false),
7255 ]);
7256 let batch = RecordBatch::try_new(
7257 Arc::new(writer_schema.clone()),
7258 vec![
7259 Arc::new(Int32Array::from(vec![i32::MIN])) as ArrayRef,
7260 Arc::new(Int64Array::from(vec![11])) as ArrayRef,
7261 ],
7262 )
7263 .expect("build writer batch");
7264 let bytes = write_ocf(&writer_schema, &[batch]);
7265 let mutated = corrupt_first_block_payload_byte(bytes, 4, 0x0f, 0x10);
7266
7267 let err = ReaderBuilder::new()
7268 .build(Cursor::new(mutated.clone()))
7269 .expect("build full reader")
7270 .collect::<Result<Vec<_>, _>>()
7271 .expect_err("full decode should reject int overflow");
7272 assert!(matches!(err, ArrowError::AvroError(_)));
7273 assert!(err.to_string().contains("varint overflow"));
7274
7275 let err = ReaderBuilder::new()
7276 .with_projection(vec![1])
7277 .build(Cursor::new(mutated))
7278 .expect("build projected reader")
7279 .collect::<Result<Vec<_>, _>>()
7280 .expect_err("projection must also reject skipped int overflow");
7281 assert!(matches!(err, ArrowError::AvroError(_)));
7282 assert!(err.to_string().contains("varint overflow"));
7283 }
7284
7285 #[test]
7286 fn comprehensive_e2e_test() {
7287 let path = "test/data/comprehensive_e2e.avro";
7288 let batch = read_file(path, 1024, false);
7289 let schema = batch.schema();
7290
7291 #[inline]
7292 fn tid_by_name(fields: &UnionFields, want: &str) -> i8 {
7293 for (tid, f) in fields.iter() {
7294 if f.name() == want {
7295 return tid;
7296 }
7297 }
7298 panic!("union child '{want}' not found");
7299 }
7300
7301 #[inline]
7302 fn tid_by_dt(fields: &UnionFields, pred: impl Fn(&DataType) -> bool) -> i8 {
7303 for (tid, f) in fields.iter() {
7304 if pred(f.data_type()) {
7305 return tid;
7306 }
7307 }
7308 panic!("no union child matches predicate");
7309 }
7310
7311 fn mk_dense_union(
7312 fields: &UnionFields,
7313 type_ids: Vec<i8>,
7314 offsets: Vec<i32>,
7315 provide: impl Fn(&Field) -> Option<ArrayRef>,
7316 ) -> ArrayRef {
7317 fn empty_child_for(dt: &DataType) -> Arc<dyn Array> {
7318 match dt {
7319 DataType::Null => Arc::new(NullArray::new(0)),
7320 DataType::Boolean => Arc::new(BooleanArray::from(Vec::<bool>::new())),
7321 DataType::Int32 => Arc::new(Int32Array::from(Vec::<i32>::new())),
7322 DataType::Int64 => Arc::new(Int64Array::from(Vec::<i64>::new())),
7323 DataType::Float32 => Arc::new(Float32Array::from(Vec::<f32>::new())),
7324 DataType::Float64 => Arc::new(Float64Array::from(Vec::<f64>::new())),
7325 DataType::Binary => Arc::new(BinaryArray::from(Vec::<&[u8]>::new())),
7326 DataType::Utf8 => Arc::new(StringArray::from(Vec::<&str>::new())),
7327 DataType::Date32 => Arc::new(Date32Array::from(Vec::<i32>::new())),
7328 DataType::Time32(arrow_schema::TimeUnit::Millisecond) => {
7329 Arc::new(Time32MillisecondArray::from(Vec::<i32>::new()))
7330 }
7331 DataType::Time64(arrow_schema::TimeUnit::Microsecond) => {
7332 Arc::new(Time64MicrosecondArray::from(Vec::<i64>::new()))
7333 }
7334 DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, tz) => {
7335 let a = TimestampMillisecondArray::from(Vec::<i64>::new());
7336 Arc::new(if let Some(tz) = tz {
7337 a.with_timezone(tz.clone())
7338 } else {
7339 a
7340 })
7341 }
7342 DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, tz) => {
7343 let a = TimestampMicrosecondArray::from(Vec::<i64>::new());
7344 Arc::new(if let Some(tz) = tz {
7345 a.with_timezone(tz.clone())
7346 } else {
7347 a
7348 })
7349 }
7350 DataType::Interval(IntervalUnit::MonthDayNano) => Arc::new(
7351 IntervalMonthDayNanoArray::from(Vec::<IntervalMonthDayNano>::new()),
7352 ),
7353 DataType::FixedSizeBinary(sz) => Arc::new(
7354 FixedSizeBinaryArray::try_from_sparse_iter_with_size(
7355 std::iter::empty::<Option<Vec<u8>>>(),
7356 *sz,
7357 )
7358 .unwrap(),
7359 ),
7360 DataType::Dictionary(_, _) => {
7361 let keys = Int32Array::from(Vec::<i32>::new());
7362 let values = Arc::new(StringArray::from(Vec::<&str>::new()));
7363 Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap())
7364 }
7365 DataType::Struct(fields) => {
7366 let children: Vec<ArrayRef> = fields
7367 .iter()
7368 .map(|f| empty_child_for(f.data_type()) as ArrayRef)
7369 .collect();
7370 Arc::new(StructArray::new(fields.clone(), children, None))
7371 }
7372 DataType::List(field) => {
7373 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0]));
7374 Arc::new(
7375 ListArray::try_new(
7376 field.clone(),
7377 offsets,
7378 empty_child_for(field.data_type()),
7379 None,
7380 )
7381 .unwrap(),
7382 )
7383 }
7384 DataType::Map(entry_field, is_sorted) => {
7385 let (key_field, val_field) = match entry_field.data_type() {
7386 DataType::Struct(fs) => (fs[0].clone(), fs[1].clone()),
7387 other => panic!("unexpected map entries type: {other:?}"),
7388 };
7389 let keys = StringArray::from(Vec::<&str>::new());
7390 let vals: ArrayRef = match val_field.data_type() {
7391 DataType::Null => Arc::new(NullArray::new(0)) as ArrayRef,
7392 DataType::Boolean => {
7393 Arc::new(BooleanArray::from(Vec::<bool>::new())) as ArrayRef
7394 }
7395 DataType::Int32 => {
7396 Arc::new(Int32Array::from(Vec::<i32>::new())) as ArrayRef
7397 }
7398 DataType::Int64 => {
7399 Arc::new(Int64Array::from(Vec::<i64>::new())) as ArrayRef
7400 }
7401 DataType::Float32 => {
7402 Arc::new(Float32Array::from(Vec::<f32>::new())) as ArrayRef
7403 }
7404 DataType::Float64 => {
7405 Arc::new(Float64Array::from(Vec::<f64>::new())) as ArrayRef
7406 }
7407 DataType::Utf8 => {
7408 Arc::new(StringArray::from(Vec::<&str>::new())) as ArrayRef
7409 }
7410 DataType::Binary => {
7411 Arc::new(BinaryArray::from(Vec::<&[u8]>::new())) as ArrayRef
7412 }
7413 DataType::Union(uf, _) => {
7414 let children: Vec<ArrayRef> = uf
7415 .iter()
7416 .map(|(_, f)| empty_child_for(f.data_type()))
7417 .collect();
7418 Arc::new(
7419 UnionArray::try_new(
7420 uf.clone(),
7421 ScalarBuffer::<i8>::from(Vec::<i8>::new()),
7422 Some(ScalarBuffer::<i32>::from(Vec::<i32>::new())),
7423 children,
7424 )
7425 .unwrap(),
7426 ) as ArrayRef
7427 }
7428 other => panic!("unsupported map value type: {other:?}"),
7429 };
7430 let entries = StructArray::new(
7431 Fields::from(vec![
7432 key_field.as_ref().clone(),
7433 val_field.as_ref().clone(),
7434 ]),
7435 vec![Arc::new(keys) as ArrayRef, vals],
7436 None,
7437 );
7438 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0]));
7439 Arc::new(MapArray::new(
7440 entry_field.clone(),
7441 offsets,
7442 entries,
7443 None,
7444 *is_sorted,
7445 ))
7446 }
7447 other => panic!("empty_child_for: unhandled type {other:?}"),
7448 }
7449 }
7450 let children: Vec<ArrayRef> = fields
7451 .iter()
7452 .map(|(_, f)| provide(f).unwrap_or_else(|| empty_child_for(f.data_type())))
7453 .collect();
7454 Arc::new(
7455 UnionArray::try_new(
7456 fields.clone(),
7457 ScalarBuffer::<i8>::from(type_ids),
7458 Some(ScalarBuffer::<i32>::from(offsets)),
7459 children,
7460 )
7461 .unwrap(),
7462 ) as ArrayRef
7463 }
7464
7465 #[inline]
7466 fn uuid16_from_str(s: &str) -> [u8; 16] {
7467 let mut out = [0u8; 16];
7468 let mut idx = 0usize;
7469 let mut hi: Option<u8> = None;
7470 for ch in s.chars() {
7471 if ch == '-' {
7472 continue;
7473 }
7474 let v = ch.to_digit(16).expect("invalid hex digit in UUID") as u8;
7475 if let Some(h) = hi {
7476 out[idx] = (h << 4) | v;
7477 idx += 1;
7478 hi = None;
7479 } else {
7480 hi = Some(v);
7481 }
7482 }
7483 assert_eq!(idx, 16, "UUID must decode to 16 bytes");
7484 out
7485 }
7486 let date_a: i32 = 19_000; let time_ms_a: i32 = 12 * 3_600_000 + 34 * 60_000 + 56_000 + 789;
7488 let time_us_eod: i64 = 86_400_000_000 - 1;
7489 let ts_ms_2024_01_01: i64 = 1_704_067_200_000; let ts_us_2024_01_01: i64 = ts_ms_2024_01_01 * 1_000;
7491 let dur_small = IntervalMonthDayNanoType::make_value(1, 2, 3_000_000_000);
7492 let dur_zero = IntervalMonthDayNanoType::make_value(0, 0, 0);
7493 let dur_large =
7494 IntervalMonthDayNanoType::make_value(12, 31, ((86_400_000 - 1) as i64) * 1_000_000);
7495 let dur_2years = IntervalMonthDayNanoType::make_value(24, 0, 0);
7496 let uuid1 = uuid16_from_str("fe7bc30b-4ce8-4c5e-b67c-2234a2d38e66");
7497 let uuid2 = uuid16_from_str("0826cc06-d2e3-4599-b4ad-af5fa6905cdb");
7498
7499 #[inline]
7500 fn push_like(
7501 reader_schema: &arrow_schema::Schema,
7502 name: &str,
7503 arr: ArrayRef,
7504 fields: &mut Vec<FieldRef>,
7505 cols: &mut Vec<ArrayRef>,
7506 ) {
7507 let src = reader_schema
7508 .field_with_name(name)
7509 .unwrap_or_else(|_| panic!("source schema missing field '{name}'"));
7510 let mut f = Field::new(name, arr.data_type().clone(), src.is_nullable());
7511 let md = src.metadata();
7512 if !md.is_empty() {
7513 f = f.with_metadata(md.clone());
7514 }
7515 fields.push(Arc::new(f));
7516 cols.push(arr);
7517 }
7518
7519 let mut fields: Vec<FieldRef> = Vec::new();
7520 let mut columns: Vec<ArrayRef> = Vec::new();
7521 push_like(
7522 schema.as_ref(),
7523 "id",
7524 Arc::new(Int64Array::from(vec![1, 2, 3, 4])) as ArrayRef,
7525 &mut fields,
7526 &mut columns,
7527 );
7528 push_like(
7529 schema.as_ref(),
7530 "flag",
7531 Arc::new(BooleanArray::from(vec![true, false, true, false])) as ArrayRef,
7532 &mut fields,
7533 &mut columns,
7534 );
7535 push_like(
7536 schema.as_ref(),
7537 "ratio_f32",
7538 Arc::new(Float32Array::from(vec![1.25f32, -0.0, 3.5, 9.75])) as ArrayRef,
7539 &mut fields,
7540 &mut columns,
7541 );
7542 push_like(
7543 schema.as_ref(),
7544 "ratio_f64",
7545 Arc::new(Float64Array::from(vec![2.5f64, -1.0, 7.0, -2.25])) as ArrayRef,
7546 &mut fields,
7547 &mut columns,
7548 );
7549 push_like(
7550 schema.as_ref(),
7551 "count_i32",
7552 Arc::new(Int32Array::from(vec![7, -1, 0, 123])) as ArrayRef,
7553 &mut fields,
7554 &mut columns,
7555 );
7556 push_like(
7557 schema.as_ref(),
7558 "count_i64",
7559 Arc::new(Int64Array::from(vec![
7560 7_000_000_000i64,
7561 -2,
7562 0,
7563 -9_876_543_210i64,
7564 ])) as ArrayRef,
7565 &mut fields,
7566 &mut columns,
7567 );
7568 push_like(
7569 schema.as_ref(),
7570 "opt_i32_nullfirst",
7571 Arc::new(Int32Array::from(vec![None, Some(42), None, Some(0)])) as ArrayRef,
7572 &mut fields,
7573 &mut columns,
7574 );
7575 push_like(
7576 schema.as_ref(),
7577 "opt_str_nullsecond",
7578 Arc::new(StringArray::from(vec![
7579 Some("alpha"),
7580 None,
7581 Some("s3"),
7582 Some(""),
7583 ])) as ArrayRef,
7584 &mut fields,
7585 &mut columns,
7586 );
7587 {
7588 let uf = match schema
7589 .field_with_name("tri_union_prim")
7590 .unwrap()
7591 .data_type()
7592 {
7593 DataType::Union(f, UnionMode::Dense) => f.clone(),
7594 other => panic!("tri_union_prim should be dense union, got {other:?}"),
7595 };
7596 let tid_i = tid_by_name(&uf, "int");
7597 let tid_s = tid_by_name(&uf, "string");
7598 let tid_b = tid_by_name(&uf, "boolean");
7599 let tids = vec![tid_i, tid_s, tid_b, tid_s];
7600 let offs = vec![0, 0, 0, 1];
7601 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
7602 DataType::Int32 => Some(Arc::new(Int32Array::from(vec![0])) as ArrayRef),
7603 DataType::Utf8 => Some(Arc::new(StringArray::from(vec!["hi", ""])) as ArrayRef),
7604 DataType::Boolean => Some(Arc::new(BooleanArray::from(vec![true])) as ArrayRef),
7605 _ => None,
7606 });
7607 push_like(
7608 schema.as_ref(),
7609 "tri_union_prim",
7610 arr,
7611 &mut fields,
7612 &mut columns,
7613 );
7614 }
7615
7616 push_like(
7617 schema.as_ref(),
7618 "str_utf8",
7619 Arc::new(StringArray::from(vec!["hello", "", "world", "✓ unicode"])) as ArrayRef,
7620 &mut fields,
7621 &mut columns,
7622 );
7623 push_like(
7624 schema.as_ref(),
7625 "raw_bytes",
7626 Arc::new(BinaryArray::from(vec![
7627 b"\x00\x01".as_ref(),
7628 b"".as_ref(),
7629 b"\xFF\x00".as_ref(),
7630 b"\x10\x20\x30\x40".as_ref(),
7631 ])) as ArrayRef,
7632 &mut fields,
7633 &mut columns,
7634 );
7635 {
7636 let it = [
7637 Some(*b"0123456789ABCDEF"),
7638 Some([0u8; 16]),
7639 Some(*b"ABCDEFGHIJKLMNOP"),
7640 Some([0xAA; 16]),
7641 ]
7642 .into_iter();
7643 let arr =
7644 Arc::new(FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 16).unwrap())
7645 as ArrayRef;
7646 push_like(
7647 schema.as_ref(),
7648 "fx16_plain",
7649 arr,
7650 &mut fields,
7651 &mut columns,
7652 );
7653 }
7654 {
7655 #[cfg(feature = "small_decimals")]
7656 let dec10_2 = Arc::new(
7657 Decimal64Array::from_iter_values([123456i64, -1, 0, 9_999_999_999i64])
7658 .with_precision_and_scale(10, 2)
7659 .unwrap(),
7660 ) as ArrayRef;
7661 #[cfg(not(feature = "small_decimals"))]
7662 let dec10_2 = Arc::new(
7663 Decimal128Array::from_iter_values([123456i128, -1, 0, 9_999_999_999i128])
7664 .with_precision_and_scale(10, 2)
7665 .unwrap(),
7666 ) as ArrayRef;
7667 push_like(
7668 schema.as_ref(),
7669 "dec_bytes_s10_2",
7670 dec10_2,
7671 &mut fields,
7672 &mut columns,
7673 );
7674 }
7675 {
7676 #[cfg(feature = "small_decimals")]
7677 let dec20_4 = Arc::new(
7678 Decimal128Array::from_iter_values([1_234_567_891_234i128, -420_000i128, 0, -1i128])
7679 .with_precision_and_scale(20, 4)
7680 .unwrap(),
7681 ) as ArrayRef;
7682 #[cfg(not(feature = "small_decimals"))]
7683 let dec20_4 = Arc::new(
7684 Decimal128Array::from_iter_values([1_234_567_891_234i128, -420_000i128, 0, -1i128])
7685 .with_precision_and_scale(20, 4)
7686 .unwrap(),
7687 ) as ArrayRef;
7688 push_like(
7689 schema.as_ref(),
7690 "dec_fix_s20_4",
7691 dec20_4,
7692 &mut fields,
7693 &mut columns,
7694 );
7695 }
7696 {
7697 let it = [Some(uuid1), Some(uuid2), Some(uuid1), Some(uuid2)].into_iter();
7698 let arr =
7699 Arc::new(FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 16).unwrap())
7700 as ArrayRef;
7701 push_like(schema.as_ref(), "uuid_str", arr, &mut fields, &mut columns);
7702 }
7703 push_like(
7704 schema.as_ref(),
7705 "d_date",
7706 Arc::new(Date32Array::from(vec![date_a, 0, 1, 365])) as ArrayRef,
7707 &mut fields,
7708 &mut columns,
7709 );
7710 push_like(
7711 schema.as_ref(),
7712 "t_millis",
7713 Arc::new(Time32MillisecondArray::from(vec![
7714 time_ms_a,
7715 0,
7716 1,
7717 86_400_000 - 1,
7718 ])) as ArrayRef,
7719 &mut fields,
7720 &mut columns,
7721 );
7722 push_like(
7723 schema.as_ref(),
7724 "t_micros",
7725 Arc::new(Time64MicrosecondArray::from(vec![
7726 time_us_eod,
7727 0,
7728 1,
7729 1_000_000,
7730 ])) as ArrayRef,
7731 &mut fields,
7732 &mut columns,
7733 );
7734 {
7735 let a = TimestampMillisecondArray::from(vec![
7736 ts_ms_2024_01_01,
7737 -1,
7738 ts_ms_2024_01_01 + 123,
7739 0,
7740 ])
7741 .with_timezone("+00:00");
7742 push_like(
7743 schema.as_ref(),
7744 "ts_millis_utc",
7745 Arc::new(a) as ArrayRef,
7746 &mut fields,
7747 &mut columns,
7748 );
7749 }
7750 {
7751 let a = TimestampMicrosecondArray::from(vec![
7752 ts_us_2024_01_01,
7753 1,
7754 ts_us_2024_01_01 + 456,
7755 0,
7756 ])
7757 .with_timezone("+00:00");
7758 push_like(
7759 schema.as_ref(),
7760 "ts_micros_utc",
7761 Arc::new(a) as ArrayRef,
7762 &mut fields,
7763 &mut columns,
7764 );
7765 }
7766 push_like(
7767 schema.as_ref(),
7768 "ts_millis_local",
7769 Arc::new(TimestampMillisecondArray::from(vec![
7770 ts_ms_2024_01_01 + 86_400_000,
7771 0,
7772 ts_ms_2024_01_01 + 789,
7773 123_456_789,
7774 ])) as ArrayRef,
7775 &mut fields,
7776 &mut columns,
7777 );
7778 push_like(
7779 schema.as_ref(),
7780 "ts_micros_local",
7781 Arc::new(TimestampMicrosecondArray::from(vec![
7782 ts_us_2024_01_01 + 123_456,
7783 0,
7784 ts_us_2024_01_01 + 101_112,
7785 987_654_321,
7786 ])) as ArrayRef,
7787 &mut fields,
7788 &mut columns,
7789 );
7790 {
7791 let v = vec![dur_small, dur_zero, dur_large, dur_2years];
7792 push_like(
7793 schema.as_ref(),
7794 "interval_mdn",
7795 Arc::new(IntervalMonthDayNanoArray::from(v)) as ArrayRef,
7796 &mut fields,
7797 &mut columns,
7798 );
7799 }
7800 {
7801 let keys = Int32Array::from(vec![1, 2, 3, 0]); let values = Arc::new(StringArray::from(vec![
7803 "UNKNOWN",
7804 "NEW",
7805 "PROCESSING",
7806 "DONE",
7807 ])) as ArrayRef;
7808 let dict = DictionaryArray::<Int32Type>::try_new(keys, values).unwrap();
7809 push_like(
7810 schema.as_ref(),
7811 "status",
7812 Arc::new(dict) as ArrayRef,
7813 &mut fields,
7814 &mut columns,
7815 );
7816 }
7817 {
7818 let list_field = match schema.field_with_name("arr_union").unwrap().data_type() {
7819 DataType::List(f) => f.clone(),
7820 other => panic!("arr_union should be List, got {other:?}"),
7821 };
7822 let uf = match list_field.data_type() {
7823 DataType::Union(f, UnionMode::Dense) => f.clone(),
7824 other => panic!("arr_union item should be union, got {other:?}"),
7825 };
7826 let tid_l = tid_by_name(&uf, "long");
7827 let tid_s = tid_by_name(&uf, "string");
7828 let tid_n = tid_by_name(&uf, "null");
7829 let type_ids = vec![
7830 tid_l, tid_s, tid_n, tid_l, tid_n, tid_s, tid_l, tid_l, tid_s, tid_n, tid_l,
7831 ];
7832 let offsets = vec![0, 0, 0, 1, 1, 1, 2, 3, 2, 2, 4];
7833 let values = mk_dense_union(&uf, type_ids, offsets, |f| match f.data_type() {
7834 DataType::Int64 => {
7835 Some(Arc::new(Int64Array::from(vec![1i64, -3, 0, -1, 0])) as ArrayRef)
7836 }
7837 DataType::Utf8 => {
7838 Some(Arc::new(StringArray::from(vec!["x", "z", "end"])) as ArrayRef)
7839 }
7840 DataType::Null => Some(Arc::new(NullArray::new(3)) as ArrayRef),
7841 _ => None,
7842 });
7843 let list_offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 4, 7, 8, 11]));
7844 let arr = Arc::new(ListArray::try_new(list_field, list_offsets, values, None).unwrap())
7845 as ArrayRef;
7846 push_like(schema.as_ref(), "arr_union", arr, &mut fields, &mut columns);
7847 }
7848 {
7849 let (entry_field, entries_fields, uf, is_sorted) =
7850 match schema.field_with_name("map_union").unwrap().data_type() {
7851 DataType::Map(entry_field, is_sorted) => {
7852 let fs = match entry_field.data_type() {
7853 DataType::Struct(fs) => fs.clone(),
7854 other => panic!("map entries must be struct, got {other:?}"),
7855 };
7856 let val_f = fs[1].clone();
7857 let uf = match val_f.data_type() {
7858 DataType::Union(f, UnionMode::Dense) => f.clone(),
7859 other => panic!("map value must be union, got {other:?}"),
7860 };
7861 (entry_field.clone(), fs, uf, *is_sorted)
7862 }
7863 other => panic!("map_union should be Map, got {other:?}"),
7864 };
7865 let keys = StringArray::from(vec!["a", "b", "c", "neg", "pi", "ok"]);
7866 let moff = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 4, 4, 6]));
7867 let tid_null = tid_by_name(&uf, "null");
7868 let tid_d = tid_by_name(&uf, "double");
7869 let tid_s = tid_by_name(&uf, "string");
7870 let type_ids = vec![tid_d, tid_null, tid_s, tid_d, tid_d, tid_s];
7871 let offsets = vec![0, 0, 0, 1, 2, 1];
7872 let pi_5dp = (std::f64::consts::PI * 100_000.0).trunc() / 100_000.0;
7873 let vals = mk_dense_union(&uf, type_ids, offsets, |f| match f.data_type() {
7874 DataType::Float64 => {
7875 Some(Arc::new(Float64Array::from(vec![1.5f64, -0.5, pi_5dp])) as ArrayRef)
7876 }
7877 DataType::Utf8 => {
7878 Some(Arc::new(StringArray::from(vec!["yes", "true"])) as ArrayRef)
7879 }
7880 DataType::Null => Some(Arc::new(NullArray::new(2)) as ArrayRef),
7881 _ => None,
7882 });
7883 let entries = StructArray::new(
7884 entries_fields.clone(),
7885 vec![Arc::new(keys) as ArrayRef, vals],
7886 None,
7887 );
7888 let map =
7889 Arc::new(MapArray::new(entry_field, moff, entries, None, is_sorted)) as ArrayRef;
7890 push_like(schema.as_ref(), "map_union", map, &mut fields, &mut columns);
7891 }
7892 {
7893 let fs = match schema.field_with_name("address").unwrap().data_type() {
7894 DataType::Struct(fs) => fs.clone(),
7895 other => panic!("address should be Struct, got {other:?}"),
7896 };
7897 let street = Arc::new(StringArray::from(vec![
7898 "100 Main",
7899 "",
7900 "42 Galaxy Way",
7901 "End Ave",
7902 ])) as ArrayRef;
7903 let zip = Arc::new(Int32Array::from(vec![12345, 0, 42424, 1])) as ArrayRef;
7904 let country = Arc::new(StringArray::from(vec!["US", "CA", "US", "GB"])) as ArrayRef;
7905 let arr = Arc::new(StructArray::new(fs, vec![street, zip, country], None)) as ArrayRef;
7906 push_like(schema.as_ref(), "address", arr, &mut fields, &mut columns);
7907 }
7908 {
7909 let fs = match schema.field_with_name("maybe_auth").unwrap().data_type() {
7910 DataType::Struct(fs) => fs.clone(),
7911 other => panic!("maybe_auth should be Struct, got {other:?}"),
7912 };
7913 let user =
7914 Arc::new(StringArray::from(vec!["alice", "bob", "carol", "dave"])) as ArrayRef;
7915 let token_values: Vec<Option<&[u8]>> = vec![
7916 None, Some(b"\x01\x02\x03".as_ref()), None, Some(b"".as_ref()), ];
7921 let token = Arc::new(BinaryArray::from(token_values)) as ArrayRef;
7922 let arr = Arc::new(StructArray::new(fs, vec![user, token], None)) as ArrayRef;
7923 push_like(
7924 schema.as_ref(),
7925 "maybe_auth",
7926 arr,
7927 &mut fields,
7928 &mut columns,
7929 );
7930 }
7931 {
7932 let uf = match schema
7933 .field_with_name("union_enum_record_array_map")
7934 .unwrap()
7935 .data_type()
7936 {
7937 DataType::Union(f, UnionMode::Dense) => f.clone(),
7938 other => panic!("union_enum_record_array_map should be union, got {other:?}"),
7939 };
7940 let mut tid_enum: Option<i8> = None;
7941 let mut tid_rec_a: Option<i8> = None;
7942 let mut tid_array: Option<i8> = None;
7943 let mut tid_map: Option<i8> = None;
7944 let mut map_entry_field: Option<FieldRef> = None;
7945 let mut map_sorted = false;
7946 for (tid, f) in uf.iter() {
7947 match f.data_type() {
7948 DataType::Dictionary(_, _) => tid_enum = Some(tid),
7949 DataType::Struct(childs)
7950 if childs.len() == 2
7951 && childs[0].name() == "a"
7952 && childs[1].name() == "b" =>
7953 {
7954 tid_rec_a = Some(tid)
7955 }
7956 DataType::List(item) if matches!(item.data_type(), DataType::Int64) => {
7957 tid_array = Some(tid)
7958 }
7959 DataType::Map(ef, is_sorted) => {
7960 tid_map = Some(tid);
7961 map_entry_field = Some(ef.clone());
7962 map_sorted = *is_sorted;
7963 }
7964 _ => {}
7965 }
7966 }
7967 let (tid_enum, tid_rec_a, tid_array, tid_map) = (
7968 tid_enum.unwrap(),
7969 tid_rec_a.unwrap(),
7970 tid_array.unwrap(),
7971 tid_map.unwrap(),
7972 );
7973 let tids = vec![tid_enum, tid_rec_a, tid_array, tid_map];
7974 let offs = vec![0, 0, 0, 0];
7975 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
7976 DataType::Dictionary(_, _) => {
7977 let keys = Int32Array::from(vec![0i32]);
7978 let values =
7979 Arc::new(StringArray::from(vec!["RED", "GREEN", "BLUE"])) as ArrayRef;
7980 Some(
7981 Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap())
7982 as ArrayRef,
7983 )
7984 }
7985 DataType::Struct(fs)
7986 if fs.len() == 2 && fs[0].name() == "a" && fs[1].name() == "b" =>
7987 {
7988 let a = Int32Array::from(vec![7]);
7989 let b = StringArray::from(vec!["rec"]);
7990 Some(Arc::new(StructArray::new(
7991 fs.clone(),
7992 vec![Arc::new(a), Arc::new(b)],
7993 None,
7994 )) as ArrayRef)
7995 }
7996 DataType::List(field) => {
7997 let values = Int64Array::from(vec![1i64, 2, 3]);
7998 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3]));
7999 Some(Arc::new(
8000 ListArray::try_new(field.clone(), offsets, Arc::new(values), None).unwrap(),
8001 ) as ArrayRef)
8002 }
8003 DataType::Map(_, _) => {
8004 let entry_field = map_entry_field.clone().unwrap();
8005 let (key_field, val_field) = match entry_field.data_type() {
8006 DataType::Struct(fs) => (fs[0].clone(), fs[1].clone()),
8007 _ => unreachable!(),
8008 };
8009 let keys = StringArray::from(vec!["k"]);
8010 let vals = StringArray::from(vec!["v"]);
8011 let entries = StructArray::new(
8012 Fields::from(vec![key_field.as_ref().clone(), val_field.as_ref().clone()]),
8013 vec![Arc::new(keys) as ArrayRef, Arc::new(vals) as ArrayRef],
8014 None,
8015 );
8016 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 1]));
8017 Some(Arc::new(MapArray::new(
8018 entry_field.clone(),
8019 offsets,
8020 entries,
8021 None,
8022 map_sorted,
8023 )) as ArrayRef)
8024 }
8025 _ => None,
8026 });
8027 push_like(
8028 schema.as_ref(),
8029 "union_enum_record_array_map",
8030 arr,
8031 &mut fields,
8032 &mut columns,
8033 );
8034 }
8035 {
8036 let uf = match schema
8037 .field_with_name("union_date_or_fixed4")
8038 .unwrap()
8039 .data_type()
8040 {
8041 DataType::Union(f, UnionMode::Dense) => f.clone(),
8042 other => panic!("union_date_or_fixed4 should be union, got {other:?}"),
8043 };
8044 let tid_date = tid_by_dt(&uf, |dt| matches!(dt, DataType::Date32));
8045 let tid_fx4 = tid_by_dt(&uf, |dt| matches!(dt, DataType::FixedSizeBinary(4)));
8046 let tids = vec![tid_date, tid_fx4, tid_date, tid_fx4];
8047 let offs = vec![0, 0, 1, 1];
8048 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
8049 DataType::Date32 => Some(Arc::new(Date32Array::from(vec![date_a, 0])) as ArrayRef),
8050 DataType::FixedSizeBinary(4) => {
8051 let it = [Some(*b"\x00\x11\x22\x33"), Some(*b"ABCD")].into_iter();
8052 Some(Arc::new(
8053 FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 4).unwrap(),
8054 ) as ArrayRef)
8055 }
8056 _ => None,
8057 });
8058 push_like(
8059 schema.as_ref(),
8060 "union_date_or_fixed4",
8061 arr,
8062 &mut fields,
8063 &mut columns,
8064 );
8065 }
8066 {
8067 let uf = match schema
8068 .field_with_name("union_interval_or_string")
8069 .unwrap()
8070 .data_type()
8071 {
8072 DataType::Union(f, UnionMode::Dense) => f.clone(),
8073 other => panic!("union_interval_or_string should be union, got {other:?}"),
8074 };
8075 let tid_dur = tid_by_dt(&uf, |dt| {
8076 matches!(dt, DataType::Interval(IntervalUnit::MonthDayNano))
8077 });
8078 let tid_str = tid_by_dt(&uf, |dt| matches!(dt, DataType::Utf8));
8079 let tids = vec![tid_dur, tid_str, tid_dur, tid_str];
8080 let offs = vec![0, 0, 1, 1];
8081 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
8082 DataType::Interval(IntervalUnit::MonthDayNano) => Some(Arc::new(
8083 IntervalMonthDayNanoArray::from(vec![dur_small, dur_large]),
8084 )
8085 as ArrayRef),
8086 DataType::Utf8 => Some(Arc::new(StringArray::from(vec![
8087 "duration-as-text",
8088 "iso-8601-period-P1Y",
8089 ])) as ArrayRef),
8090 _ => None,
8091 });
8092 push_like(
8093 schema.as_ref(),
8094 "union_interval_or_string",
8095 arr,
8096 &mut fields,
8097 &mut columns,
8098 );
8099 }
8100 {
8101 let uf = match schema
8102 .field_with_name("union_uuid_or_fixed10")
8103 .unwrap()
8104 .data_type()
8105 {
8106 DataType::Union(f, UnionMode::Dense) => f.clone(),
8107 other => panic!("union_uuid_or_fixed10 should be union, got {other:?}"),
8108 };
8109 let tid_uuid = tid_by_dt(&uf, |dt| matches!(dt, DataType::FixedSizeBinary(16)));
8110 let tid_fx10 = tid_by_dt(&uf, |dt| matches!(dt, DataType::FixedSizeBinary(10)));
8111 let tids = vec![tid_uuid, tid_fx10, tid_uuid, tid_fx10];
8112 let offs = vec![0, 0, 1, 1];
8113 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
8114 DataType::FixedSizeBinary(16) => {
8115 let it = [Some(uuid1), Some(uuid2)].into_iter();
8116 Some(Arc::new(
8117 FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 16).unwrap(),
8118 ) as ArrayRef)
8119 }
8120 DataType::FixedSizeBinary(10) => {
8121 let fx10_a = [0xAAu8; 10];
8122 let fx10_b = [0x00u8, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99];
8123 let it = [Some(fx10_a), Some(fx10_b)].into_iter();
8124 Some(Arc::new(
8125 FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 10).unwrap(),
8126 ) as ArrayRef)
8127 }
8128 _ => None,
8129 });
8130 push_like(
8131 schema.as_ref(),
8132 "union_uuid_or_fixed10",
8133 arr,
8134 &mut fields,
8135 &mut columns,
8136 );
8137 }
8138 {
8139 let list_field = match schema
8140 .field_with_name("array_records_with_union")
8141 .unwrap()
8142 .data_type()
8143 {
8144 DataType::List(f) => f.clone(),
8145 other => panic!("array_records_with_union should be List, got {other:?}"),
8146 };
8147 let kv_fields = match list_field.data_type() {
8148 DataType::Struct(fs) => fs.clone(),
8149 other => panic!("array_records_with_union items must be Struct, got {other:?}"),
8150 };
8151 let val_field = kv_fields
8152 .iter()
8153 .find(|f| f.name() == "val")
8154 .unwrap()
8155 .clone();
8156 let uf = match val_field.data_type() {
8157 DataType::Union(f, UnionMode::Dense) => f.clone(),
8158 other => panic!("KV.val should be union, got {other:?}"),
8159 };
8160 let keys = Arc::new(StringArray::from(vec!["k1", "k2", "k", "k3", "x"])) as ArrayRef;
8161 let tid_null = tid_by_name(&uf, "null");
8162 let tid_i = tid_by_name(&uf, "int");
8163 let tid_l = tid_by_name(&uf, "long");
8164 let type_ids = vec![tid_i, tid_null, tid_l, tid_null, tid_i];
8165 let offsets = vec![0, 0, 0, 1, 1];
8166 let vals = mk_dense_union(&uf, type_ids, offsets, |f| match f.data_type() {
8167 DataType::Int32 => Some(Arc::new(Int32Array::from(vec![5, -5])) as ArrayRef),
8168 DataType::Int64 => Some(Arc::new(Int64Array::from(vec![99i64])) as ArrayRef),
8169 DataType::Null => Some(Arc::new(NullArray::new(2)) as ArrayRef),
8170 _ => None,
8171 });
8172 let values_struct =
8173 Arc::new(StructArray::new(kv_fields.clone(), vec![keys, vals], None)) as ArrayRef;
8174 let list_offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2, 3, 4, 5]));
8175 let arr = Arc::new(
8176 ListArray::try_new(list_field, list_offsets, values_struct, None).unwrap(),
8177 ) as ArrayRef;
8178 push_like(
8179 schema.as_ref(),
8180 "array_records_with_union",
8181 arr,
8182 &mut fields,
8183 &mut columns,
8184 );
8185 }
8186 {
8187 let uf = match schema
8188 .field_with_name("union_map_or_array_int")
8189 .unwrap()
8190 .data_type()
8191 {
8192 DataType::Union(f, UnionMode::Dense) => f.clone(),
8193 other => panic!("union_map_or_array_int should be union, got {other:?}"),
8194 };
8195 let tid_map = tid_by_dt(&uf, |dt| matches!(dt, DataType::Map(_, _)));
8196 let tid_list = tid_by_dt(&uf, |dt| matches!(dt, DataType::List(_)));
8197 let map_child: ArrayRef = {
8198 let (entry_field, is_sorted) = match uf
8199 .iter()
8200 .find(|(tid, _)| *tid == tid_map)
8201 .unwrap()
8202 .1
8203 .data_type()
8204 {
8205 DataType::Map(ef, is_sorted) => (ef.clone(), *is_sorted),
8206 _ => unreachable!(),
8207 };
8208 let (key_field, val_field) = match entry_field.data_type() {
8209 DataType::Struct(fs) => (fs[0].clone(), fs[1].clone()),
8210 _ => unreachable!(),
8211 };
8212 let keys = StringArray::from(vec!["x", "y", "only"]);
8213 let vals = Int32Array::from(vec![1, 2, 10]);
8214 let entries = StructArray::new(
8215 Fields::from(vec![key_field.as_ref().clone(), val_field.as_ref().clone()]),
8216 vec![Arc::new(keys) as ArrayRef, Arc::new(vals) as ArrayRef],
8217 None,
8218 );
8219 let moff = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2, 3]));
8220 Arc::new(MapArray::new(entry_field, moff, entries, None, is_sorted)) as ArrayRef
8221 };
8222 let list_child: ArrayRef = {
8223 let list_field = match uf
8224 .iter()
8225 .find(|(tid, _)| *tid == tid_list)
8226 .unwrap()
8227 .1
8228 .data_type()
8229 {
8230 DataType::List(f) => f.clone(),
8231 _ => unreachable!(),
8232 };
8233 let values = Int32Array::from(vec![1, 2, 3, 0]);
8234 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 4]));
8235 Arc::new(ListArray::try_new(list_field, offsets, Arc::new(values), None).unwrap())
8236 as ArrayRef
8237 };
8238 let tids = vec![tid_map, tid_list, tid_map, tid_list];
8239 let offs = vec![0, 0, 1, 1];
8240 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
8241 DataType::Map(_, _) => Some(map_child.clone()),
8242 DataType::List(_) => Some(list_child.clone()),
8243 _ => None,
8244 });
8245 push_like(
8246 schema.as_ref(),
8247 "union_map_or_array_int",
8248 arr,
8249 &mut fields,
8250 &mut columns,
8251 );
8252 }
8253 push_like(
8254 schema.as_ref(),
8255 "renamed_with_default",
8256 Arc::new(Int32Array::from(vec![100, 42, 7, 42])) as ArrayRef,
8257 &mut fields,
8258 &mut columns,
8259 );
8260 {
8261 let fs = match schema.field_with_name("person").unwrap().data_type() {
8262 DataType::Struct(fs) => fs.clone(),
8263 other => panic!("person should be Struct, got {other:?}"),
8264 };
8265 let name =
8266 Arc::new(StringArray::from(vec!["Alice", "Bob", "Carol", "Dave"])) as ArrayRef;
8267 let age = Arc::new(Int32Array::from(vec![30, 0, 25, 41])) as ArrayRef;
8268 let arr = Arc::new(StructArray::new(fs, vec![name, age], None)) as ArrayRef;
8269 push_like(schema.as_ref(), "person", arr, &mut fields, &mut columns);
8270 }
8271 let expected =
8272 RecordBatch::try_new(Arc::new(Schema::new(Fields::from(fields))), columns).unwrap();
8273 assert_eq!(
8274 expected, batch,
8275 "entire RecordBatch mismatch (schema, all columns, all rows)"
8276 );
8277 }
8278 #[test]
8279 #[cfg_attr(miri, ignore)] fn comprehensive_e2e_resolution_test() {
8281 use serde_json::Value;
8282 use std::collections::HashMap;
8283
8284 fn make_comprehensive_reader_schema(path: &str) -> AvroSchema {
8297 fn set_type_string(f: &mut Value, new_ty: &str) {
8298 if let Some(ty) = f.get_mut("type") {
8299 match ty {
8300 Value::String(_) | Value::Object(_) => {
8301 *ty = Value::String(new_ty.to_string());
8302 }
8303 Value::Array(arr) => {
8304 for b in arr.iter_mut() {
8305 match b {
8306 Value::String(s) if s != "null" => {
8307 *b = Value::String(new_ty.to_string());
8308 break;
8309 }
8310 Value::Object(_) => {
8311 *b = Value::String(new_ty.to_string());
8312 break;
8313 }
8314 _ => {}
8315 }
8316 }
8317 }
8318 _ => {}
8319 }
8320 }
8321 }
8322 fn reverse_union_array(f: &mut Value) {
8323 if let Some(arr) = f.get_mut("type").and_then(|t| t.as_array_mut()) {
8324 arr.reverse();
8325 }
8326 }
8327 fn reverse_items_union(f: &mut Value) {
8328 if let Some(obj) = f.get_mut("type").and_then(|t| t.as_object_mut())
8329 && let Some(items) = obj.get_mut("items").and_then(|v| v.as_array_mut())
8330 {
8331 items.reverse();
8332 }
8333 }
8334 fn reverse_map_values_union(f: &mut Value) {
8335 if let Some(obj) = f.get_mut("type").and_then(|t| t.as_object_mut())
8336 && let Some(values) = obj.get_mut("values").and_then(|v| v.as_array_mut())
8337 {
8338 values.reverse();
8339 }
8340 }
8341 fn reverse_nested_union_in_record(f: &mut Value, field_name: &str) {
8342 if let Some(obj) = f.get_mut("type").and_then(|t| t.as_object_mut())
8343 && let Some(fields) = obj.get_mut("fields").and_then(|v| v.as_array_mut())
8344 {
8345 for ff in fields.iter_mut() {
8346 if ff.get("name").and_then(|n| n.as_str()) == Some(field_name)
8347 && let Some(ty) = ff.get_mut("type")
8348 && let Some(arr) = ty.as_array_mut()
8349 {
8350 arr.reverse();
8351 }
8352 }
8353 }
8354 }
8355 fn rename_nested_field_with_alias(f: &mut Value, old: &str, new: &str) {
8356 if let Some(obj) = f.get_mut("type").and_then(|t| t.as_object_mut())
8357 && let Some(fields) = obj.get_mut("fields").and_then(|v| v.as_array_mut())
8358 {
8359 for ff in fields.iter_mut() {
8360 if ff.get("name").and_then(|n| n.as_str()) == Some(old) {
8361 ff["name"] = Value::String(new.to_string());
8362 ff["aliases"] = Value::Array(vec![Value::String(old.to_string())]);
8363 }
8364 }
8365 }
8366 }
8367 let mut root = load_writer_schema_json(path);
8368 assert_eq!(root["type"], "record", "writer schema must be a record");
8369 let fields = root
8370 .get_mut("fields")
8371 .and_then(|f| f.as_array_mut())
8372 .expect("record has fields");
8373 for f in fields.iter_mut() {
8374 let Some(name) = f.get("name").and_then(|n| n.as_str()) else {
8375 continue;
8376 };
8377 match name {
8378 "id" => {
8380 f["name"] = Value::String("identifier".into());
8381 f["aliases"] = Value::Array(vec![Value::String("id".into())]);
8382 }
8383 "renamed_with_default" => {
8384 f["name"] = Value::String("old_count".into());
8385 f["aliases"] =
8386 Value::Array(vec![Value::String("renamed_with_default".into())]);
8387 }
8388 "count_i32" => set_type_string(f, "long"),
8390 "ratio_f32" => set_type_string(f, "double"),
8391 "opt_str_nullsecond" => reverse_union_array(f),
8393 "union_enum_record_array_map" => reverse_union_array(f),
8394 "union_date_or_fixed4" => reverse_union_array(f),
8395 "union_interval_or_string" => reverse_union_array(f),
8396 "union_uuid_or_fixed10" => reverse_union_array(f),
8397 "union_map_or_array_int" => reverse_union_array(f),
8398 "maybe_auth" => reverse_nested_union_in_record(f, "token"),
8399 "arr_union" => reverse_items_union(f),
8401 "map_union" => reverse_map_values_union(f),
8402 "address" => rename_nested_field_with_alias(f, "street", "street_name"),
8404 "person" => {
8406 if let Some(tobj) = f.get_mut("type").and_then(|t| t.as_object_mut()) {
8407 tobj.insert("name".to_string(), Value::String("Person".into()));
8408 tobj.insert(
8409 "namespace".to_string(),
8410 Value::String("com.example".into()),
8411 );
8412 tobj.insert(
8413 "aliases".into(),
8414 Value::Array(vec![
8415 Value::String("PersonV2".into()),
8416 Value::String("com.example.v2.PersonV2".into()),
8417 ]),
8418 );
8419 }
8420 }
8421 _ => {}
8422 }
8423 }
8424 fields.reverse();
8425 AvroSchema::new(root.to_string())
8426 }
8427
8428 let path = "test/data/comprehensive_e2e.avro";
8429 let reader_schema = make_comprehensive_reader_schema(path);
8430 let batch = read_alltypes_with_reader_schema(path, reader_schema.clone());
8431
8432 const UUID_EXT_KEY: &str = "ARROW:extension:name";
8433 const UUID_LOGICAL_KEY: &str = "logicalType";
8434
8435 let uuid_md_top: Option<arrow_schema::Metadata> = batch
8436 .schema()
8437 .field_with_name("uuid_str")
8438 .ok()
8439 .and_then(|f| {
8440 let md = f.metadata();
8441 let has_ext = md.get(UUID_EXT_KEY).is_some();
8442 let is_uuid_logical = md
8443 .get(UUID_LOGICAL_KEY)
8444 .map(|v| v.trim_matches('"') == "uuid")
8445 .unwrap_or(false);
8446 if has_ext || is_uuid_logical {
8447 Some(md.clone())
8448 } else {
8449 None
8450 }
8451 });
8452
8453 let uuid_md_union: Option<arrow_schema::Metadata> = batch
8454 .schema()
8455 .field_with_name("union_uuid_or_fixed10")
8456 .ok()
8457 .and_then(|f| match f.data_type() {
8458 DataType::Union(uf, _) => {
8459 let (_, child) = uf.iter().find(|(_, child)| child.name() == "uuid")?;
8460 let md = child.metadata();
8461 let has_ext = md.get(UUID_EXT_KEY).is_some();
8462 let is_uuid_logical = md
8463 .get(UUID_LOGICAL_KEY)
8464 .map(|v| v.trim_matches('"') == "uuid")
8465 .unwrap_or(false);
8466 if has_ext || is_uuid_logical {
8467 Some(md.clone())
8468 } else {
8469 None
8470 }
8471 }
8472 _ => None,
8473 });
8474
8475 let add_uuid_ext_top = |f: Field| -> Field {
8476 if let Some(md) = &uuid_md_top {
8477 f.with_metadata(md.clone())
8478 } else {
8479 f
8480 }
8481 };
8482 let add_uuid_ext_union = |f: Field| -> Field {
8483 if let Some(md) = &uuid_md_union {
8484 f.with_metadata(md.clone())
8485 } else {
8486 f
8487 }
8488 };
8489
8490 #[inline]
8491 fn uuid16_from_str(s: &str) -> [u8; 16] {
8492 let mut out = [0u8; 16];
8493 let mut idx = 0usize;
8494 let mut hi: Option<u8> = None;
8495 for ch in s.chars() {
8496 if ch == '-' {
8497 continue;
8498 }
8499 let v = ch.to_digit(16).expect("invalid hex digit in UUID") as u8;
8500 if let Some(h) = hi {
8501 out[idx] = (h << 4) | v;
8502 idx += 1;
8503 hi = None;
8504 } else {
8505 hi = Some(v);
8506 }
8507 }
8508 assert_eq!(idx, 16, "UUID must decode to 16 bytes");
8509 out
8510 }
8511
8512 fn mk_dense_union(
8513 fields: &UnionFields,
8514 type_ids: Vec<i8>,
8515 offsets: Vec<i32>,
8516 provide: impl Fn(&Field) -> Option<ArrayRef>,
8517 ) -> ArrayRef {
8518 fn empty_child_for(dt: &DataType) -> Arc<dyn Array> {
8519 match dt {
8520 DataType::Null => Arc::new(NullArray::new(0)),
8521 DataType::Boolean => Arc::new(BooleanArray::from(Vec::<bool>::new())),
8522 DataType::Int32 => Arc::new(Int32Array::from(Vec::<i32>::new())),
8523 DataType::Int64 => Arc::new(Int64Array::from(Vec::<i64>::new())),
8524 DataType::Float32 => Arc::new(Float32Array::from(Vec::<f32>::new())),
8525 DataType::Float64 => Arc::new(Float64Array::from(Vec::<f64>::new())),
8526 DataType::Binary => Arc::new(BinaryArray::from(Vec::<&[u8]>::new())),
8527 DataType::Utf8 => Arc::new(StringArray::from(Vec::<&str>::new())),
8528 DataType::Date32 => Arc::new(Date32Array::from(Vec::<i32>::new())),
8529 DataType::Time32(arrow_schema::TimeUnit::Millisecond) => {
8530 Arc::new(Time32MillisecondArray::from(Vec::<i32>::new()))
8531 }
8532 DataType::Time64(arrow_schema::TimeUnit::Microsecond) => {
8533 Arc::new(Time64MicrosecondArray::from(Vec::<i64>::new()))
8534 }
8535 DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, tz) => {
8536 let a = TimestampMillisecondArray::from(Vec::<i64>::new());
8537 Arc::new(if let Some(tz) = tz {
8538 a.with_timezone(tz.clone())
8539 } else {
8540 a
8541 })
8542 }
8543 DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, tz) => {
8544 let a = TimestampMicrosecondArray::from(Vec::<i64>::new());
8545 Arc::new(if let Some(tz) = tz {
8546 a.with_timezone(tz.clone())
8547 } else {
8548 a
8549 })
8550 }
8551 DataType::Interval(IntervalUnit::MonthDayNano) => Arc::new(
8552 IntervalMonthDayNanoArray::from(Vec::<IntervalMonthDayNano>::new()),
8553 ),
8554 DataType::FixedSizeBinary(sz) => Arc::new(
8555 FixedSizeBinaryArray::try_from_sparse_iter_with_size(
8556 std::iter::empty::<Option<Vec<u8>>>(),
8557 *sz,
8558 )
8559 .unwrap(),
8560 ),
8561 DataType::Dictionary(_, _) => {
8562 let keys = Int32Array::from(Vec::<i32>::new());
8563 let values = Arc::new(StringArray::from(Vec::<&str>::new()));
8564 Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap())
8565 }
8566 DataType::Struct(fields) => {
8567 let children: Vec<ArrayRef> = fields
8568 .iter()
8569 .map(|f| empty_child_for(f.data_type()) as ArrayRef)
8570 .collect();
8571 Arc::new(StructArray::new(fields.clone(), children, None))
8572 }
8573 DataType::List(field) => {
8574 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0]));
8575 Arc::new(
8576 ListArray::try_new(
8577 field.clone(),
8578 offsets,
8579 empty_child_for(field.data_type()),
8580 None,
8581 )
8582 .unwrap(),
8583 )
8584 }
8585 DataType::Map(entry_field, is_sorted) => {
8586 let (key_field, val_field) = match entry_field.data_type() {
8587 DataType::Struct(fs) => (fs[0].clone(), fs[1].clone()),
8588 other => panic!("unexpected map entries type: {other:?}"),
8589 };
8590 let keys = StringArray::from(Vec::<&str>::new());
8591 let vals: ArrayRef = match val_field.data_type() {
8592 DataType::Null => Arc::new(NullArray::new(0)) as ArrayRef,
8593 DataType::Boolean => {
8594 Arc::new(BooleanArray::from(Vec::<bool>::new())) as ArrayRef
8595 }
8596 DataType::Int32 => {
8597 Arc::new(Int32Array::from(Vec::<i32>::new())) as ArrayRef
8598 }
8599 DataType::Int64 => {
8600 Arc::new(Int64Array::from(Vec::<i64>::new())) as ArrayRef
8601 }
8602 DataType::Float32 => {
8603 Arc::new(Float32Array::from(Vec::<f32>::new())) as ArrayRef
8604 }
8605 DataType::Float64 => {
8606 Arc::new(Float64Array::from(Vec::<f64>::new())) as ArrayRef
8607 }
8608 DataType::Utf8 => {
8609 Arc::new(StringArray::from(Vec::<&str>::new())) as ArrayRef
8610 }
8611 DataType::Binary => {
8612 Arc::new(BinaryArray::from(Vec::<&[u8]>::new())) as ArrayRef
8613 }
8614 DataType::Union(uf, _) => {
8615 let children: Vec<ArrayRef> = uf
8616 .iter()
8617 .map(|(_, f)| empty_child_for(f.data_type()))
8618 .collect();
8619 Arc::new(
8620 UnionArray::try_new(
8621 uf.clone(),
8622 ScalarBuffer::<i8>::from(Vec::<i8>::new()),
8623 Some(ScalarBuffer::<i32>::from(Vec::<i32>::new())),
8624 children,
8625 )
8626 .unwrap(),
8627 ) as ArrayRef
8628 }
8629 other => panic!("unsupported map value type: {other:?}"),
8630 };
8631 let entries = StructArray::new(
8632 Fields::from(vec![
8633 key_field.as_ref().clone(),
8634 val_field.as_ref().clone(),
8635 ]),
8636 vec![Arc::new(keys) as ArrayRef, vals],
8637 None,
8638 );
8639 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0]));
8640 Arc::new(MapArray::new(
8641 entry_field.clone(),
8642 offsets,
8643 entries,
8644 None,
8645 *is_sorted,
8646 ))
8647 }
8648 other => panic!("empty_child_for: unhandled type {other:?}"),
8649 }
8650 }
8651 let children: Vec<ArrayRef> = fields
8652 .iter()
8653 .map(|(_, f)| provide(f).unwrap_or_else(|| empty_child_for(f.data_type())))
8654 .collect();
8655 Arc::new(
8656 UnionArray::try_new(
8657 fields.clone(),
8658 ScalarBuffer::<i8>::from(type_ids),
8659 Some(ScalarBuffer::<i32>::from(offsets)),
8660 children,
8661 )
8662 .unwrap(),
8663 ) as ArrayRef
8664 }
8665 let date_a: i32 = 19_000; let time_ms_a: i32 = 12 * 3_600_000 + 34 * 60_000 + 56_000 + 789;
8667 let time_us_eod: i64 = 86_400_000_000 - 1;
8668 let ts_ms_2024_01_01: i64 = 1_704_067_200_000; let ts_us_2024_01_01: i64 = ts_ms_2024_01_01 * 1_000;
8670 let dur_small = IntervalMonthDayNanoType::make_value(1, 2, 3_000_000_000);
8671 let dur_zero = IntervalMonthDayNanoType::make_value(0, 0, 0);
8672 let dur_large =
8673 IntervalMonthDayNanoType::make_value(12, 31, ((86_400_000 - 1) as i64) * 1_000_000);
8674 let dur_2years = IntervalMonthDayNanoType::make_value(24, 0, 0);
8675 let uuid1 = uuid16_from_str("fe7bc30b-4ce8-4c5e-b67c-2234a2d38e66");
8676 let uuid2 = uuid16_from_str("0826cc06-d2e3-4599-b4ad-af5fa6905cdb");
8677 let item_name = Field::LIST_FIELD_DEFAULT_NAME;
8678 let uf_tri = UnionFields::try_new(
8679 vec![0, 1, 2],
8680 vec![
8681 Field::new("int", DataType::Int32, false),
8682 Field::new("string", DataType::Utf8, false),
8683 Field::new("boolean", DataType::Boolean, false),
8684 ],
8685 )
8686 .unwrap();
8687 let uf_arr_items = UnionFields::try_new(
8688 vec![0, 1, 2],
8689 vec![
8690 Field::new("null", DataType::Null, false),
8691 Field::new("string", DataType::Utf8, false),
8692 Field::new("long", DataType::Int64, false),
8693 ],
8694 )
8695 .unwrap();
8696 let arr_items_field = Arc::new(Field::new(
8697 item_name,
8698 DataType::Union(uf_arr_items.clone(), UnionMode::Dense),
8699 true,
8700 ));
8701 let uf_map_vals = UnionFields::try_new(
8702 vec![0, 1, 2],
8703 vec![
8704 Field::new("string", DataType::Utf8, false),
8705 Field::new("double", DataType::Float64, false),
8706 Field::new("null", DataType::Null, false),
8707 ],
8708 )
8709 .unwrap();
8710 let map_entries_field = Arc::new(Field::new(
8711 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
8712 DataType::Struct(Fields::from(vec![
8713 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
8714 Field::new(
8715 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
8716 DataType::Union(uf_map_vals.clone(), UnionMode::Dense),
8717 true,
8718 ),
8719 ])),
8720 false,
8721 ));
8722 let mut enum_md_color = {
8724 let mut m = HashMap::<String, String>::new();
8725 m.insert(
8726 crate::schema::AVRO_ENUM_SYMBOLS_METADATA_KEY.to_string(),
8727 serde_json::to_string(&vec!["RED", "GREEN", "BLUE"]).unwrap(),
8728 );
8729 m
8730 };
8731 enum_md_color.insert(AVRO_NAME_METADATA_KEY.to_string(), "Color".to_string());
8732 enum_md_color.insert(
8733 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8734 "org.apache.arrow.avrotests.v1.types".to_string(),
8735 );
8736 let union_rec_a_fields = Fields::from(vec![
8737 Field::new("a", DataType::Int32, false),
8738 Field::new("b", DataType::Utf8, false),
8739 ]);
8740 let union_rec_b_fields = Fields::from(vec![
8741 Field::new("x", DataType::Int64, false),
8742 Field::new("y", DataType::Binary, false),
8743 ]);
8744 let union_map_entries = Arc::new(Field::new(
8745 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
8746 DataType::Struct(Fields::from(vec![
8747 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
8748 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, false),
8749 ])),
8750 false,
8751 ));
8752 let person_md = {
8753 let mut m = HashMap::<String, String>::new();
8754 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "Person".to_string());
8755 m.insert(
8756 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8757 "com.example".to_string(),
8758 );
8759 m
8760 };
8761 let maybe_auth_md = {
8762 let mut m = HashMap::<String, String>::new();
8763 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "MaybeAuth".to_string());
8764 m.insert(
8765 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8766 "org.apache.arrow.avrotests.v1.types".to_string(),
8767 );
8768 m
8769 };
8770 let address_md = {
8771 let mut m = HashMap::<String, String>::new();
8772 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "Address".to_string());
8773 m.insert(
8774 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8775 "org.apache.arrow.avrotests.v1.types".to_string(),
8776 );
8777 m
8778 };
8779 let rec_a_md = {
8780 let mut m = HashMap::<String, String>::new();
8781 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "RecA".to_string());
8782 m.insert(
8783 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8784 "org.apache.arrow.avrotests.v1.types".to_string(),
8785 );
8786 m
8787 };
8788 let rec_b_md = {
8789 let mut m = HashMap::<String, String>::new();
8790 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "RecB".to_string());
8791 m.insert(
8792 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8793 "org.apache.arrow.avrotests.v1.types".to_string(),
8794 );
8795 m
8796 };
8797 let uf_union_big = UnionFields::try_new(
8798 vec![0, 1, 2, 3, 4],
8799 vec![
8800 Field::new(
8801 "map",
8802 DataType::Map(union_map_entries.clone(), false),
8803 false,
8804 ),
8805 Field::new(
8806 "array",
8807 DataType::List(Arc::new(Field::new(item_name, DataType::Int64, false))),
8808 false,
8809 ),
8810 Field::new(
8811 "org.apache.arrow.avrotests.v1.types.RecB",
8812 DataType::Struct(union_rec_b_fields.clone()),
8813 false,
8814 )
8815 .with_metadata(rec_b_md.clone()),
8816 Field::new(
8817 "org.apache.arrow.avrotests.v1.types.RecA",
8818 DataType::Struct(union_rec_a_fields.clone()),
8819 false,
8820 )
8821 .with_metadata(rec_a_md.clone()),
8822 Field::new(
8823 "org.apache.arrow.avrotests.v1.types.Color",
8824 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
8825 false,
8826 )
8827 .with_metadata(enum_md_color.clone()),
8828 ],
8829 )
8830 .unwrap();
8831 let fx4_md = {
8832 let mut m = HashMap::<String, String>::new();
8833 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "Fx4".to_string());
8834 m.insert(
8835 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8836 "org.apache.arrow.avrotests.v1".to_string(),
8837 );
8838 m
8839 };
8840 let uf_date_fixed4 = UnionFields::try_new(
8841 vec![0, 1],
8842 vec![
8843 Field::new(
8844 "org.apache.arrow.avrotests.v1.Fx4",
8845 DataType::FixedSizeBinary(4),
8846 false,
8847 )
8848 .with_metadata(fx4_md.clone()),
8849 Field::new("date", DataType::Date32, false),
8850 ],
8851 )
8852 .unwrap();
8853 let dur12u_md = {
8854 let mut m = HashMap::<String, String>::new();
8855 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "Dur12U".to_string());
8856 m.insert(
8857 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8858 "org.apache.arrow.avrotests.v1".to_string(),
8859 );
8860 m
8861 };
8862 let uf_dur_or_str = UnionFields::try_new(
8863 vec![0, 1],
8864 vec![
8865 Field::new("string", DataType::Utf8, false),
8866 Field::new(
8867 "org.apache.arrow.avrotests.v1.Dur12U",
8868 DataType::Interval(arrow_schema::IntervalUnit::MonthDayNano),
8869 false,
8870 )
8871 .with_metadata(dur12u_md.clone()),
8872 ],
8873 )
8874 .unwrap();
8875 let fx10_md = {
8876 let mut m = HashMap::<String, String>::new();
8877 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "Fx10".to_string());
8878 m.insert(
8879 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8880 "org.apache.arrow.avrotests.v1".to_string(),
8881 );
8882 m
8883 };
8884 let uf_uuid_or_fx10 = UnionFields::try_new(
8885 vec![0, 1],
8886 vec![
8887 Field::new(
8888 "org.apache.arrow.avrotests.v1.Fx10",
8889 DataType::FixedSizeBinary(10),
8890 false,
8891 )
8892 .with_metadata(fx10_md.clone()),
8893 add_uuid_ext_union(Field::new("uuid", DataType::FixedSizeBinary(16), false)),
8894 ],
8895 )
8896 .unwrap();
8897 let uf_kv_val = UnionFields::try_new(
8898 vec![0, 1, 2],
8899 vec![
8900 Field::new("null", DataType::Null, false),
8901 Field::new("int", DataType::Int32, false),
8902 Field::new("long", DataType::Int64, false),
8903 ],
8904 )
8905 .unwrap();
8906 let kv_fields = Fields::from(vec![
8907 Field::new("key", DataType::Utf8, false),
8908 Field::new(
8909 "val",
8910 DataType::Union(uf_kv_val.clone(), UnionMode::Dense),
8911 true,
8912 ),
8913 ]);
8914 let kv_md = {
8915 let mut m = HashMap::<String, String>::new();
8916 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "KV".to_string());
8917 m.insert(
8918 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8919 "org.apache.arrow.avrotests.v1.types".to_string(),
8920 );
8921 m
8922 };
8923 let kv_item_field = Arc::new(
8924 Field::new(item_name, DataType::Struct(kv_fields.clone()), false).with_metadata(kv_md),
8925 );
8926 let map_int_entries = Arc::new(Field::new(
8927 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
8928 DataType::Struct(Fields::from(vec![
8929 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
8930 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Int32, false),
8931 ])),
8932 false,
8933 ));
8934 let uf_map_or_array = UnionFields::try_new(
8935 vec![0, 1],
8936 vec![
8937 Field::new(
8938 "array",
8939 DataType::List(Arc::new(Field::new(item_name, DataType::Int32, false))),
8940 false,
8941 ),
8942 Field::new("map", DataType::Map(map_int_entries.clone(), false), false),
8943 ],
8944 )
8945 .unwrap();
8946 let mut enum_md_status = {
8947 let mut m = HashMap::<String, String>::new();
8948 m.insert(
8949 crate::schema::AVRO_ENUM_SYMBOLS_METADATA_KEY.to_string(),
8950 serde_json::to_string(&vec!["UNKNOWN", "NEW", "PROCESSING", "DONE"]).unwrap(),
8951 );
8952 m
8953 };
8954 enum_md_status.insert(AVRO_NAME_METADATA_KEY.to_string(), "Status".to_string());
8955 enum_md_status.insert(
8956 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8957 "org.apache.arrow.avrotests.v1.types".to_string(),
8958 );
8959 let mut dec20_md = HashMap::<String, String>::new();
8960 dec20_md.insert("precision".to_string(), "20".to_string());
8961 dec20_md.insert("scale".to_string(), "4".to_string());
8962 dec20_md.insert(AVRO_NAME_METADATA_KEY.to_string(), "DecFix20".to_string());
8963 dec20_md.insert(
8964 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8965 "org.apache.arrow.avrotests.v1.types".to_string(),
8966 );
8967 let mut dec10_md = HashMap::<String, String>::new();
8968 dec10_md.insert("precision".to_string(), "10".to_string());
8969 dec10_md.insert("scale".to_string(), "2".to_string());
8970 let fx16_top_md = {
8971 let mut m = HashMap::<String, String>::new();
8972 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "Fx16".to_string());
8973 m.insert(
8974 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8975 "org.apache.arrow.avrotests.v1.types".to_string(),
8976 );
8977 m
8978 };
8979 let dur12_top_md = {
8980 let mut m = HashMap::<String, String>::new();
8981 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "Dur12".to_string());
8982 m.insert(
8983 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8984 "org.apache.arrow.avrotests.v1.types".to_string(),
8985 );
8986 m
8987 };
8988 #[cfg(feature = "small_decimals")]
8989 let dec20_dt = DataType::Decimal128(20, 4);
8990 #[cfg(not(feature = "small_decimals"))]
8991 let dec20_dt = DataType::Decimal128(20, 4);
8992 #[cfg(feature = "small_decimals")]
8993 let dec10_dt = DataType::Decimal64(10, 2);
8994 #[cfg(not(feature = "small_decimals"))]
8995 let dec10_dt = DataType::Decimal128(10, 2);
8996 let fields: Vec<FieldRef> = vec![
8997 Arc::new(
8998 Field::new(
8999 "person",
9000 DataType::Struct(Fields::from(vec![
9001 Field::new("name", DataType::Utf8, false),
9002 Field::new("age", DataType::Int32, false),
9003 ])),
9004 false,
9005 )
9006 .with_metadata(person_md),
9007 ),
9008 Arc::new(Field::new("old_count", DataType::Int32, false)),
9009 Arc::new(Field::new(
9010 "union_map_or_array_int",
9011 DataType::Union(uf_map_or_array.clone(), UnionMode::Dense),
9012 false,
9013 )),
9014 Arc::new(Field::new(
9015 "array_records_with_union",
9016 DataType::List(kv_item_field.clone()),
9017 false,
9018 )),
9019 Arc::new(Field::new(
9020 "union_uuid_or_fixed10",
9021 DataType::Union(uf_uuid_or_fx10.clone(), UnionMode::Dense),
9022 false,
9023 )),
9024 Arc::new(Field::new(
9025 "union_interval_or_string",
9026 DataType::Union(uf_dur_or_str.clone(), UnionMode::Dense),
9027 false,
9028 )),
9029 Arc::new(Field::new(
9030 "union_date_or_fixed4",
9031 DataType::Union(uf_date_fixed4.clone(), UnionMode::Dense),
9032 false,
9033 )),
9034 Arc::new(Field::new(
9035 "union_enum_record_array_map",
9036 DataType::Union(uf_union_big.clone(), UnionMode::Dense),
9037 false,
9038 )),
9039 Arc::new(
9040 Field::new(
9041 "maybe_auth",
9042 DataType::Struct(Fields::from(vec![
9043 Field::new("user", DataType::Utf8, false),
9044 Field::new("token", DataType::Binary, true), ])),
9046 false,
9047 )
9048 .with_metadata(maybe_auth_md),
9049 ),
9050 Arc::new(
9051 Field::new(
9052 "address",
9053 DataType::Struct(Fields::from(vec![
9054 Field::new("street_name", DataType::Utf8, false),
9055 Field::new("zip", DataType::Int32, false),
9056 Field::new("country", DataType::Utf8, false),
9057 ])),
9058 false,
9059 )
9060 .with_metadata(address_md),
9061 ),
9062 Arc::new(Field::new(
9063 "map_union",
9064 DataType::Map(map_entries_field.clone(), false),
9065 false,
9066 )),
9067 Arc::new(Field::new(
9068 "arr_union",
9069 DataType::List(arr_items_field.clone()),
9070 false,
9071 )),
9072 Arc::new(
9073 Field::new(
9074 "status",
9075 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
9076 false,
9077 )
9078 .with_metadata(enum_md_status.clone()),
9079 ),
9080 Arc::new(
9081 Field::new(
9082 "interval_mdn",
9083 DataType::Interval(IntervalUnit::MonthDayNano),
9084 false,
9085 )
9086 .with_metadata(dur12_top_md.clone()),
9087 ),
9088 Arc::new(Field::new(
9089 "ts_micros_local",
9090 DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, None),
9091 false,
9092 )),
9093 Arc::new(Field::new(
9094 "ts_millis_local",
9095 DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, None),
9096 false,
9097 )),
9098 Arc::new(Field::new(
9099 "ts_micros_utc",
9100 DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, Some("+00:00".into())),
9101 false,
9102 )),
9103 Arc::new(Field::new(
9104 "ts_millis_utc",
9105 DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, Some("+00:00".into())),
9106 false,
9107 )),
9108 Arc::new(Field::new(
9109 "t_micros",
9110 DataType::Time64(arrow_schema::TimeUnit::Microsecond),
9111 false,
9112 )),
9113 Arc::new(Field::new(
9114 "t_millis",
9115 DataType::Time32(arrow_schema::TimeUnit::Millisecond),
9116 false,
9117 )),
9118 Arc::new(Field::new("d_date", DataType::Date32, false)),
9119 Arc::new(add_uuid_ext_top(Field::new(
9120 "uuid_str",
9121 DataType::FixedSizeBinary(16),
9122 false,
9123 ))),
9124 Arc::new(Field::new("dec_fix_s20_4", dec20_dt, false).with_metadata(dec20_md.clone())),
9125 Arc::new(
9126 Field::new("dec_bytes_s10_2", dec10_dt, false).with_metadata(dec10_md.clone()),
9127 ),
9128 Arc::new(
9129 Field::new("fx16_plain", DataType::FixedSizeBinary(16), false)
9130 .with_metadata(fx16_top_md.clone()),
9131 ),
9132 Arc::new(Field::new("raw_bytes", DataType::Binary, false)),
9133 Arc::new(Field::new("str_utf8", DataType::Utf8, false)),
9134 Arc::new(Field::new(
9135 "tri_union_prim",
9136 DataType::Union(uf_tri.clone(), UnionMode::Dense),
9137 false,
9138 )),
9139 Arc::new(Field::new("opt_str_nullsecond", DataType::Utf8, true)),
9140 Arc::new(Field::new("opt_i32_nullfirst", DataType::Int32, true)),
9141 Arc::new(Field::new("count_i64", DataType::Int64, false)),
9142 Arc::new(Field::new("count_i32", DataType::Int64, false)),
9143 Arc::new(Field::new("ratio_f64", DataType::Float64, false)),
9144 Arc::new(Field::new("ratio_f32", DataType::Float64, false)),
9145 Arc::new(Field::new("flag", DataType::Boolean, false)),
9146 Arc::new(Field::new("identifier", DataType::Int64, false)),
9147 ];
9148 let expected_schema = Arc::new(arrow_schema::Schema::new(Fields::from(fields)));
9149 let mut cols: Vec<ArrayRef> = vec![
9150 Arc::new(StructArray::new(
9151 match expected_schema
9152 .field_with_name("person")
9153 .unwrap()
9154 .data_type()
9155 {
9156 DataType::Struct(fs) => fs.clone(),
9157 _ => unreachable!(),
9158 },
9159 vec![
9160 Arc::new(StringArray::from(vec!["Alice", "Bob", "Carol", "Dave"])) as ArrayRef,
9161 Arc::new(Int32Array::from(vec![30, 0, 25, 41])) as ArrayRef,
9162 ],
9163 None,
9164 )) as ArrayRef,
9165 Arc::new(Int32Array::from(vec![100, 42, 7, 42])) as ArrayRef,
9166 ];
9167 {
9168 let map_child: ArrayRef = {
9169 let keys = StringArray::from(vec!["x", "y", "only"]);
9170 let vals = Int32Array::from(vec![1, 2, 10]);
9171 let entries = StructArray::new(
9172 Fields::from(vec![
9173 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
9174 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Int32, false),
9175 ]),
9176 vec![Arc::new(keys) as ArrayRef, Arc::new(vals) as ArrayRef],
9177 None,
9178 );
9179 let moff = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2, 3]));
9180 Arc::new(MapArray::new(
9181 map_int_entries.clone(),
9182 moff,
9183 entries,
9184 None,
9185 false,
9186 )) as ArrayRef
9187 };
9188 let list_child: ArrayRef = {
9189 let values = Int32Array::from(vec![1, 2, 3, 0]);
9190 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 4]));
9191 Arc::new(
9192 ListArray::try_new(
9193 Arc::new(Field::new(item_name, DataType::Int32, false)),
9194 offsets,
9195 Arc::new(values),
9196 None,
9197 )
9198 .unwrap(),
9199 ) as ArrayRef
9200 };
9201 let tids = vec![1, 0, 1, 0];
9202 let offs = vec![0, 0, 1, 1];
9203 let arr = mk_dense_union(&uf_map_or_array, tids, offs, |f| match f.name().as_str() {
9204 "array" => Some(list_child.clone()),
9205 "map" => Some(map_child.clone()),
9206 _ => None,
9207 });
9208 cols.push(arr);
9209 }
9210 {
9211 let keys = Arc::new(StringArray::from(vec!["k1", "k2", "k", "k3", "x"])) as ArrayRef;
9212 let type_ids = vec![1, 0, 2, 0, 1];
9213 let offsets = vec![0, 0, 0, 1, 1];
9214 let vals = mk_dense_union(&uf_kv_val, type_ids, offsets, |f| match f.data_type() {
9215 DataType::Int32 => Some(Arc::new(Int32Array::from(vec![5, -5])) as ArrayRef),
9216 DataType::Int64 => Some(Arc::new(Int64Array::from(vec![99i64])) as ArrayRef),
9217 DataType::Null => Some(Arc::new(NullArray::new(2)) as ArrayRef),
9218 _ => None,
9219 });
9220 let values_struct =
9221 Arc::new(StructArray::new(kv_fields.clone(), vec![keys, vals], None));
9222 let list_offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2, 3, 4, 5]));
9223 let arr = Arc::new(
9224 ListArray::try_new(kv_item_field.clone(), list_offsets, values_struct, None)
9225 .unwrap(),
9226 ) as ArrayRef;
9227 cols.push(arr);
9228 }
9229 {
9230 let type_ids = vec![1, 0, 1, 0]; let offs = vec![0, 0, 1, 1];
9232 let arr = mk_dense_union(&uf_uuid_or_fx10, type_ids, offs, |f| match f.data_type() {
9233 DataType::FixedSizeBinary(16) => {
9234 let it = [Some(uuid1), Some(uuid2)].into_iter();
9235 Some(Arc::new(
9236 FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 16).unwrap(),
9237 ) as ArrayRef)
9238 }
9239 DataType::FixedSizeBinary(10) => {
9240 let fx10_a = [0xAAu8; 10];
9241 let fx10_b = [0x00u8, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99];
9242 let it = [Some(fx10_a), Some(fx10_b)].into_iter();
9243 Some(Arc::new(
9244 FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 10).unwrap(),
9245 ) as ArrayRef)
9246 }
9247 _ => None,
9248 });
9249 cols.push(arr);
9250 }
9251 {
9252 let type_ids = vec![1, 0, 1, 0]; let offs = vec![0, 0, 1, 1];
9254 let arr = mk_dense_union(&uf_dur_or_str, type_ids, offs, |f| match f.data_type() {
9255 DataType::Interval(arrow_schema::IntervalUnit::MonthDayNano) => Some(Arc::new(
9256 IntervalMonthDayNanoArray::from(vec![dur_small, dur_large]),
9257 )
9258 as ArrayRef),
9259 DataType::Utf8 => Some(Arc::new(StringArray::from(vec![
9260 "duration-as-text",
9261 "iso-8601-period-P1Y",
9262 ])) as ArrayRef),
9263 _ => None,
9264 });
9265 cols.push(arr);
9266 }
9267 {
9268 let type_ids = vec![1, 0, 1, 0]; let offs = vec![0, 0, 1, 1];
9270 let arr = mk_dense_union(&uf_date_fixed4, type_ids, offs, |f| match f.data_type() {
9271 DataType::Date32 => Some(Arc::new(Date32Array::from(vec![date_a, 0])) as ArrayRef),
9272 DataType::FixedSizeBinary(4) => {
9273 let it = [Some(*b"\x00\x11\x22\x33"), Some(*b"ABCD")].into_iter();
9274 Some(Arc::new(
9275 FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 4).unwrap(),
9276 ) as ArrayRef)
9277 }
9278 _ => None,
9279 });
9280 cols.push(arr);
9281 }
9282 {
9283 let tids = vec![4, 3, 1, 0]; let offs = vec![0, 0, 0, 0];
9285 let arr = mk_dense_union(&uf_union_big, tids, offs, |f| match f.data_type() {
9286 DataType::Dictionary(_, _) => {
9287 let keys = Int32Array::from(vec![0i32]);
9288 let values =
9289 Arc::new(StringArray::from(vec!["RED", "GREEN", "BLUE"])) as ArrayRef;
9290 Some(
9291 Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap())
9292 as ArrayRef,
9293 )
9294 }
9295 DataType::Struct(fs) if fs == &union_rec_a_fields => {
9296 let a = Int32Array::from(vec![7]);
9297 let b = StringArray::from(vec!["rec"]);
9298 Some(Arc::new(StructArray::new(
9299 fs.clone(),
9300 vec![Arc::new(a) as ArrayRef, Arc::new(b) as ArrayRef],
9301 None,
9302 )) as ArrayRef)
9303 }
9304 DataType::List(_) => {
9305 let values = Int64Array::from(vec![1i64, 2, 3]);
9306 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3]));
9307 Some(Arc::new(
9308 ListArray::try_new(
9309 Arc::new(Field::new(item_name, DataType::Int64, false)),
9310 offsets,
9311 Arc::new(values),
9312 None,
9313 )
9314 .unwrap(),
9315 ) as ArrayRef)
9316 }
9317 DataType::Map(_, _) => {
9318 let keys = StringArray::from(vec!["k"]);
9319 let vals = StringArray::from(vec!["v"]);
9320 let entries = StructArray::new(
9321 Fields::from(vec![
9322 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
9323 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, false),
9324 ]),
9325 vec![Arc::new(keys) as ArrayRef, Arc::new(vals) as ArrayRef],
9326 None,
9327 );
9328 let moff = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 1]));
9329 Some(Arc::new(MapArray::new(
9330 union_map_entries.clone(),
9331 moff,
9332 entries,
9333 None,
9334 false,
9335 )) as ArrayRef)
9336 }
9337 _ => None,
9338 });
9339 cols.push(arr);
9340 }
9341 {
9342 let fs = match expected_schema
9343 .field_with_name("maybe_auth")
9344 .unwrap()
9345 .data_type()
9346 {
9347 DataType::Struct(fs) => fs.clone(),
9348 _ => unreachable!(),
9349 };
9350 let user =
9351 Arc::new(StringArray::from(vec!["alice", "bob", "carol", "dave"])) as ArrayRef;
9352 let token_values: Vec<Option<&[u8]>> = vec![
9353 None,
9354 Some(b"\x01\x02\x03".as_ref()),
9355 None,
9356 Some(b"".as_ref()),
9357 ];
9358 let token = Arc::new(BinaryArray::from(token_values)) as ArrayRef;
9359 cols.push(Arc::new(StructArray::new(fs, vec![user, token], None)) as ArrayRef);
9360 }
9361 {
9362 let fs = match expected_schema
9363 .field_with_name("address")
9364 .unwrap()
9365 .data_type()
9366 {
9367 DataType::Struct(fs) => fs.clone(),
9368 _ => unreachable!(),
9369 };
9370 let street = Arc::new(StringArray::from(vec![
9371 "100 Main",
9372 "",
9373 "42 Galaxy Way",
9374 "End Ave",
9375 ])) as ArrayRef;
9376 let zip = Arc::new(Int32Array::from(vec![12345, 0, 42424, 1])) as ArrayRef;
9377 let country = Arc::new(StringArray::from(vec!["US", "CA", "US", "GB"])) as ArrayRef;
9378 cols.push(Arc::new(StructArray::new(fs, vec![street, zip, country], None)) as ArrayRef);
9379 }
9380 {
9381 let keys = StringArray::from(vec!["a", "b", "c", "neg", "pi", "ok"]);
9382 let moff = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 4, 4, 6]));
9383 let tid_s = 0; let tid_d = 1; let tid_n = 2; let type_ids = vec![tid_d, tid_n, tid_s, tid_d, tid_d, tid_s];
9387 let offsets = vec![0, 0, 0, 1, 2, 1];
9388 let pi_5dp = (std::f64::consts::PI * 100_000.0).trunc() / 100_000.0;
9389 let vals = mk_dense_union(&uf_map_vals, type_ids, offsets, |f| match f.data_type() {
9390 DataType::Float64 => {
9391 Some(Arc::new(Float64Array::from(vec![1.5f64, -0.5, pi_5dp])) as ArrayRef)
9392 }
9393 DataType::Utf8 => {
9394 Some(Arc::new(StringArray::from(vec!["yes", "true"])) as ArrayRef)
9395 }
9396 DataType::Null => Some(Arc::new(NullArray::new(1)) as ArrayRef),
9397 _ => None,
9398 });
9399 let entries = StructArray::new(
9400 Fields::from(vec![
9401 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
9402 Field::new(
9403 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
9404 DataType::Union(uf_map_vals.clone(), UnionMode::Dense),
9405 true,
9406 ),
9407 ]),
9408 vec![Arc::new(keys) as ArrayRef, vals],
9409 None,
9410 );
9411 let map = Arc::new(MapArray::new(
9412 map_entries_field.clone(),
9413 moff,
9414 entries,
9415 None,
9416 false,
9417 )) as ArrayRef;
9418 cols.push(map);
9419 }
9420 {
9421 let type_ids = vec![
9422 2, 1, 0, 2, 0, 1, 2, 2, 1, 0,
9423 2, ];
9425 let offsets = vec![0, 0, 0, 1, 1, 1, 2, 3, 2, 2, 4];
9426 let values =
9427 mk_dense_union(&uf_arr_items, type_ids, offsets, |f| match f.data_type() {
9428 DataType::Int64 => {
9429 Some(Arc::new(Int64Array::from(vec![1i64, -3, 0, -1, 0])) as ArrayRef)
9430 }
9431 DataType::Utf8 => {
9432 Some(Arc::new(StringArray::from(vec!["x", "z", "end"])) as ArrayRef)
9433 }
9434 DataType::Null => Some(Arc::new(NullArray::new(3)) as ArrayRef),
9435 _ => None,
9436 });
9437 let list_offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 4, 7, 8, 11]));
9438 let arr = Arc::new(
9439 ListArray::try_new(arr_items_field.clone(), list_offsets, values, None).unwrap(),
9440 ) as ArrayRef;
9441 cols.push(arr);
9442 }
9443 {
9444 let keys = Int32Array::from(vec![1, 2, 3, 0]); let values = Arc::new(StringArray::from(vec![
9446 "UNKNOWN",
9447 "NEW",
9448 "PROCESSING",
9449 "DONE",
9450 ])) as ArrayRef;
9451 let dict = DictionaryArray::<Int32Type>::try_new(keys, values).unwrap();
9452 cols.push(Arc::new(dict) as ArrayRef);
9453 }
9454 cols.push(Arc::new(IntervalMonthDayNanoArray::from(vec![
9455 dur_small, dur_zero, dur_large, dur_2years,
9456 ])) as ArrayRef);
9457 cols.push(Arc::new(TimestampMicrosecondArray::from(vec![
9458 ts_us_2024_01_01 + 123_456,
9459 0,
9460 ts_us_2024_01_01 + 101_112,
9461 987_654_321,
9462 ])) as ArrayRef);
9463 cols.push(Arc::new(TimestampMillisecondArray::from(vec![
9464 ts_ms_2024_01_01 + 86_400_000,
9465 0,
9466 ts_ms_2024_01_01 + 789,
9467 123_456_789,
9468 ])) as ArrayRef);
9469 {
9470 let a = TimestampMicrosecondArray::from(vec![
9471 ts_us_2024_01_01,
9472 1,
9473 ts_us_2024_01_01 + 456,
9474 0,
9475 ])
9476 .with_timezone("+00:00");
9477 cols.push(Arc::new(a) as ArrayRef);
9478 }
9479 {
9480 let a = TimestampMillisecondArray::from(vec![
9481 ts_ms_2024_01_01,
9482 -1,
9483 ts_ms_2024_01_01 + 123,
9484 0,
9485 ])
9486 .with_timezone("+00:00");
9487 cols.push(Arc::new(a) as ArrayRef);
9488 }
9489 cols.push(Arc::new(Time64MicrosecondArray::from(vec![
9490 time_us_eod,
9491 0,
9492 1,
9493 1_000_000,
9494 ])) as ArrayRef);
9495 cols.push(Arc::new(Time32MillisecondArray::from(vec![
9496 time_ms_a,
9497 0,
9498 1,
9499 86_400_000 - 1,
9500 ])) as ArrayRef);
9501 cols.push(Arc::new(Date32Array::from(vec![date_a, 0, 1, 365])) as ArrayRef);
9502 {
9503 let it = [Some(uuid1), Some(uuid2), Some(uuid1), Some(uuid2)].into_iter();
9504 cols.push(Arc::new(
9505 FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 16).unwrap(),
9506 ) as ArrayRef);
9507 }
9508 {
9509 #[cfg(feature = "small_decimals")]
9510 let arr = Arc::new(
9511 Decimal128Array::from_iter_values([1_234_567_891_234i128, -420_000i128, 0, -1i128])
9512 .with_precision_and_scale(20, 4)
9513 .unwrap(),
9514 ) as ArrayRef;
9515 #[cfg(not(feature = "small_decimals"))]
9516 let arr = Arc::new(
9517 Decimal128Array::from_iter_values([1_234_567_891_234i128, -420_000i128, 0, -1i128])
9518 .with_precision_and_scale(20, 4)
9519 .unwrap(),
9520 ) as ArrayRef;
9521 cols.push(arr);
9522 }
9523 {
9524 #[cfg(feature = "small_decimals")]
9525 let arr = Arc::new(
9526 Decimal64Array::from_iter_values([123456i64, -1, 0, 9_999_999_999i64])
9527 .with_precision_and_scale(10, 2)
9528 .unwrap(),
9529 ) as ArrayRef;
9530 #[cfg(not(feature = "small_decimals"))]
9531 let arr = Arc::new(
9532 Decimal128Array::from_iter_values([123456i128, -1, 0, 9_999_999_999i128])
9533 .with_precision_and_scale(10, 2)
9534 .unwrap(),
9535 ) as ArrayRef;
9536 cols.push(arr);
9537 }
9538 {
9539 let it = [
9540 Some(*b"0123456789ABCDEF"),
9541 Some([0u8; 16]),
9542 Some(*b"ABCDEFGHIJKLMNOP"),
9543 Some([0xAA; 16]),
9544 ]
9545 .into_iter();
9546 cols.push(Arc::new(
9547 FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 16).unwrap(),
9548 ) as ArrayRef);
9549 }
9550 cols.push(Arc::new(BinaryArray::from(vec![
9551 b"\x00\x01".as_ref(),
9552 b"".as_ref(),
9553 b"\xFF\x00".as_ref(),
9554 b"\x10\x20\x30\x40".as_ref(),
9555 ])) as ArrayRef);
9556 cols.push(Arc::new(StringArray::from(vec!["hello", "", "world", "✓ unicode"])) as ArrayRef);
9557 {
9558 let tids = vec![0, 1, 2, 1];
9559 let offs = vec![0, 0, 0, 1];
9560 let arr = mk_dense_union(&uf_tri, tids, offs, |f| match f.data_type() {
9561 DataType::Int32 => Some(Arc::new(Int32Array::from(vec![0])) as ArrayRef),
9562 DataType::Utf8 => Some(Arc::new(StringArray::from(vec!["hi", ""])) as ArrayRef),
9563 DataType::Boolean => Some(Arc::new(BooleanArray::from(vec![true])) as ArrayRef),
9564 _ => None,
9565 });
9566 cols.push(arr);
9567 }
9568 cols.push(Arc::new(StringArray::from(vec![
9569 Some("alpha"),
9570 None,
9571 Some("s3"),
9572 Some(""),
9573 ])) as ArrayRef);
9574 cols.push(Arc::new(Int32Array::from(vec![None, Some(42), None, Some(0)])) as ArrayRef);
9575 cols.push(Arc::new(Int64Array::from(vec![
9576 7_000_000_000i64,
9577 -2,
9578 0,
9579 -9_876_543_210i64,
9580 ])) as ArrayRef);
9581 cols.push(Arc::new(Int64Array::from(vec![7i64, -1, 0, 123])) as ArrayRef);
9582 cols.push(Arc::new(Float64Array::from(vec![2.5f64, -1.0, 7.0, -2.25])) as ArrayRef);
9583 cols.push(Arc::new(Float64Array::from(vec![1.25f64, -0.0, 3.5, 9.75])) as ArrayRef);
9584 cols.push(Arc::new(BooleanArray::from(vec![true, false, true, false])) as ArrayRef);
9585 cols.push(Arc::new(Int64Array::from(vec![1, 2, 3, 4])) as ArrayRef);
9586 let expected = RecordBatch::try_new(expected_schema, cols).unwrap();
9587 assert_eq!(
9588 expected, batch,
9589 "entire RecordBatch mismatch (schema, all columns, all rows)"
9590 );
9591 }
9592
9593 fn make_type_ref_ocf() -> Vec<u8> {
9603 use apache_avro::{Schema as ApacheSchema, Writer as ApacheWriter, types::Value};
9604 let schema_json = r#"{
9605 "type": "record", "name": "Root",
9606 "fields": [
9607 {"name": "ts", "type": {"type": "record", "name": "Timestamp", "fields": [
9608 {"name": "seconds", "type": "long"},
9609 {"name": "nanos", "type": "int"}
9610 ]}},
9611 {"name": "extra", "type": {"type": "record", "name": "Event", "fields": [
9612 {"name": "time", "type": "Timestamp"}
9613 ]}}
9614 ]
9615 }"#;
9616 let schema = ApacheSchema::parse_str(schema_json).expect("valid schema");
9617 let mut out = Vec::new();
9618 {
9619 let mut writer = ApacheWriter::new(&schema, &mut out).unwrap();
9620 let ts_val = |s: i64, n: i32| {
9621 Value::Record(vec![
9622 ("seconds".into(), Value::Long(s)),
9623 ("nanos".into(), Value::Int(n)),
9624 ])
9625 };
9626 for (ts_s, ts_n, ex_s, ex_n) in [(1000i64, 100i32, -1i64, -1i32), (2000, 200, -2, -2)] {
9628 let row = Value::Record(vec![
9629 ("ts".into(), ts_val(ts_s, ts_n)),
9630 (
9631 "extra".into(),
9632 Value::Record(vec![("time".into(), ts_val(ex_s, ex_n))]),
9633 ),
9634 ]);
9635 writer.append_value_ref(&row).expect("append row");
9636 }
9637 writer.flush().expect("flush");
9638 }
9639 out
9640 }
9641
9642 #[test]
9651 fn test_nullable_reader_schema_vs_plain_writer_nested_struct() {
9652 let bytes = make_type_ref_ocf();
9653 let reader_schema = AvroSchema::new(
9654 r#"{"type":"record","name":"Root","fields":[
9655 {"name":"ts","type":["null",{"type":"record","name":"Timestamp","fields":[
9656 {"name":"seconds","type":["null","long"]},
9657 {"name":"nanos", "type":["null","int"]}
9658 ]}]}
9659 ]}"#
9660 .to_string(),
9661 );
9662 let mut reader = ReaderBuilder::new()
9663 .with_reader_schema(reader_schema)
9664 .build(Cursor::new(bytes))
9665 .expect("reader should build");
9666 let batch = reader
9667 .next()
9668 .expect("should have a batch")
9669 .expect("reading should succeed");
9670 assert_eq!(batch.num_rows(), 2);
9671 let ts = batch
9672 .column(0)
9673 .as_any()
9674 .downcast_ref::<StructArray>()
9675 .unwrap();
9676 let seconds = ts
9677 .column_by_name("seconds")
9678 .unwrap()
9679 .as_any()
9680 .downcast_ref::<Int64Array>()
9681 .unwrap();
9682 assert_eq!(seconds.value(0), 1000);
9683 assert_eq!(seconds.value(1), 2000);
9684 }
9685
9686 #[test]
9694 fn test_skipper_consumes_writer_only_struct_fields() {
9695 let bytes = make_type_ref_ocf();
9696 let reader_schema = AvroSchema::new(
9697 r#"{"type":"record","name":"Root","fields":[
9698 {"name":"ts","type":{"type":"record","name":"Timestamp","fields":[
9699 {"name":"seconds","type":"long"}
9700 ]}}
9701 ]}"#
9702 .to_string(),
9703 );
9704 let mut reader = ReaderBuilder::new()
9705 .with_reader_schema(reader_schema)
9706 .build(Cursor::new(bytes))
9707 .expect("reader should build");
9708 let batch = reader
9709 .next()
9710 .expect("should have a batch")
9711 .expect("Skipper must consume both seconds and nanos for extra.time");
9712 assert_eq!(batch.num_rows(), 2);
9713 let ts = batch
9714 .column(0)
9715 .as_any()
9716 .downcast_ref::<StructArray>()
9717 .unwrap();
9718 let seconds = ts
9719 .column_by_name("seconds")
9720 .unwrap()
9721 .as_any()
9722 .downcast_ref::<Int64Array>()
9723 .unwrap();
9724 assert_eq!(seconds.value(0), 1000);
9725 assert_eq!(seconds.value(1), 2000);
9726 }
9727
9728 #[test]
9737 fn test_skip_array_of_structs_uses_writer_schema_not_resolved() {
9738 use apache_avro::{Schema as ApacheSchema, Writer as ApacheWriter, types::Value};
9739 let schema_json = r#"{
9740 "type": "record", "name": "Root",
9741 "fields": [
9742 {"name": "ts", "type": {"type": "record", "name": "Timestamp", "fields": [
9743 {"name": "seconds", "type": "long"},
9744 {"name": "nanos", "type": "int"}
9745 ]}},
9746 {"name": "events", "type": {"type": "array", "items": {
9747 "type": "record", "name": "Event", "fields": [
9748 {"name": "time", "type": "Timestamp"}
9749 ]
9750 }}}
9751 ]
9752 }"#;
9753 let schema = ApacheSchema::parse_str(schema_json).expect("valid schema");
9754 let mut bytes = Vec::new();
9755 {
9756 let mut writer = ApacheWriter::new(&schema, &mut bytes).unwrap();
9757 let ts_val = |s: i64, n: i32| {
9759 Value::Record(vec![
9760 ("seconds".into(), Value::Long(s)),
9761 ("nanos".into(), Value::Int(n)),
9762 ])
9763 };
9764 let row = Value::Record(vec![
9765 ("ts".into(), ts_val(100, 5)),
9766 (
9767 "events".into(),
9768 Value::Array(vec![Value::Record(vec![("time".into(), ts_val(200, 1))])]),
9769 ),
9770 ]);
9771 writer.append_value_ref(&row).expect("append row");
9772 writer.flush().expect("flush");
9773 }
9774
9775 let reader_schema = AvroSchema::new(
9777 r#"{"type":"record","name":"Root","fields":[
9778 {"name":"ts","type":["null",{"type":"record","name":"Timestamp","fields":[
9779 {"name":"seconds","type":["null","long"]},
9780 {"name":"nanos", "type":["null","int"]}
9781 ]}]}
9782 ]}"#
9783 .to_string(),
9784 );
9785 let mut reader = ReaderBuilder::new()
9786 .with_reader_schema(reader_schema)
9787 .build(Cursor::new(bytes))
9788 .expect("reader should build");
9789 let batch = reader
9790 .next()
9791 .expect("should have a batch")
9792 .expect("Skipper must consume all events bytes using writer field types");
9793 assert_eq!(batch.num_rows(), 1);
9794 let ts = batch
9795 .column(0)
9796 .as_any()
9797 .downcast_ref::<StructArray>()
9798 .unwrap();
9799 let seconds = ts
9800 .column_by_name("seconds")
9801 .unwrap()
9802 .as_any()
9803 .downcast_ref::<Int64Array>()
9804 .unwrap();
9805 assert_eq!(seconds.value(0), 100);
9806 }
9807}