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#[allow(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 avro_schema = match store.lookup(&fingerprint) {
1104 Some(schema) => schema,
1105 None => {
1106 return Err(AvroError::General(format!(
1107 "Fingerprint {fingerprint:?} not found in schema store",
1108 )));
1109 }
1110 };
1111 let writer_schema = avro_schema.schema()?;
1112 let record_decoder = match projection {
1113 None => self.make_record_decoder_from_schemas(&writer_schema, reader_schema)?,
1114 Some(projection) => {
1115 if let Some(ref pruned_reader_schema) = projected_reader_schema {
1116 self.make_record_decoder_from_schemas(
1117 &writer_schema,
1118 Some(pruned_reader_schema),
1119 )?
1120 } else {
1121 let derived_reader_schema = avro_schema.project(projection)?;
1122 self.make_record_decoder_from_schemas(
1123 &writer_schema,
1124 Some(&derived_reader_schema),
1125 )?
1126 }
1127 }
1128 };
1129 if fingerprint == start_fingerprint {
1130 active_decoder = Some(record_decoder);
1131 } else {
1132 cache.insert(fingerprint, record_decoder);
1133 }
1134 }
1135 let active_decoder = active_decoder.ok_or_else(|| {
1136 AvroError::General(format!(
1137 "Initial fingerprint {start_fingerprint:?} not found in schema store"
1138 ))
1139 })?;
1140 Ok(Decoder::from_parts(
1141 self.batch_size,
1142 active_decoder,
1143 Some(start_fingerprint),
1144 cache,
1145 store.fingerprint_algorithm(),
1146 ))
1147 }
1148
1149 pub fn with_batch_size(mut self, batch_size: usize) -> Self {
1155 self.batch_size = batch_size;
1156 self
1157 }
1158
1159 pub fn with_utf8_view(mut self, utf8_view: bool) -> Self {
1165 self.utf8_view = utf8_view;
1166 self
1167 }
1168
1169 pub fn use_utf8view(&self) -> bool {
1171 self.utf8_view
1172 }
1173
1174 pub fn with_strict_mode(mut self, strict_mode: bool) -> Self {
1179 self.strict_mode = strict_mode;
1180 self
1181 }
1182
1183 pub fn with_tz(mut self, tz: Tz) -> Self {
1187 self.tz = tz;
1188 self
1189 }
1190
1191 pub fn with_reader_schema(mut self, schema: AvroSchema) -> Self {
1198 self.reader_schema = Some(schema);
1199 self
1200 }
1201
1202 pub fn with_projection(mut self, projection: Vec<usize>) -> Self {
1260 self.projection = Some(projection);
1261 self
1262 }
1263
1264 pub fn with_writer_schema_store(mut self, store: SchemaStore) -> Self {
1272 self.writer_schema_store = Some(store);
1273 self
1274 }
1275
1276 pub fn with_active_fingerprint(mut self, fp: Fingerprint) -> Self {
1281 self.active_fingerprint = Some(fp);
1282 self
1283 }
1284
1285 pub fn build<R: BufRead>(self, mut reader: R) -> Result<Reader<R>, ArrowError> {
1291 let (header, _) = read_header(&mut reader)?;
1292 let decoder = self.make_decoder(Some(&header), self.reader_schema.as_ref())?;
1293 Ok(Reader {
1294 reader,
1295 header,
1296 decoder,
1297 block_decoder: BlockDecoder::default(),
1298 block_data: Vec::new(),
1299 block_count: 0,
1300 block_cursor: 0,
1301 finished: false,
1302 })
1303 }
1304
1305 pub fn build_decoder(self) -> Result<Decoder, ArrowError> {
1314 if self.writer_schema_store.is_none() {
1315 return Err(ArrowError::InvalidArgumentError(
1316 "Building a decoder requires a writer schema store".to_string(),
1317 ));
1318 }
1319 self.make_decoder(None, self.reader_schema.as_ref())
1320 .map_err(ArrowError::from)
1321 }
1322}
1323
1324#[derive(Debug)]
1334pub struct Reader<R: BufRead> {
1335 reader: R,
1336 header: Header,
1337 decoder: Decoder,
1338 block_decoder: BlockDecoder,
1339 block_data: Vec<u8>,
1340 block_count: usize,
1341 block_cursor: usize,
1342 finished: bool,
1343}
1344
1345impl<R: BufRead> Reader<R> {
1346 pub fn schema(&self) -> SchemaRef {
1349 self.decoder.schema()
1350 }
1351
1352 pub fn avro_header(&self) -> &Header {
1354 &self.header
1355 }
1356
1357 fn read(&mut self) -> Result<Option<RecordBatch>, AvroError> {
1362 'outer: while !self.finished && !self.decoder.batch_is_full() {
1363 while self.block_cursor == self.block_data.len() {
1364 let buf = self.reader.fill_buf()?;
1365 if buf.is_empty() {
1366 self.finished = true;
1367 break 'outer;
1368 }
1369 let consumed = self.block_decoder.decode(buf)?;
1371 self.reader.consume(consumed);
1372 if let Some(block) = self.block_decoder.flush() {
1373 if block.sync != self.header.sync() {
1375 return Err(AvroError::ParseError(
1376 "Avro block sync marker does not match file header".to_string(),
1377 ));
1378 }
1379 self.block_data = if let Some(ref codec) = self.header.compression()? {
1380 let decompressed: Vec<u8> = codec.decompress(&block.data)?;
1381 decompressed
1382 } else {
1383 block.data
1384 };
1385 self.block_count = block.count;
1386 self.block_cursor = 0;
1387 } else if consumed == 0 {
1388 return Err(AvroError::ParseError(
1390 "Could not decode next Avro block from partial data".to_string(),
1391 ));
1392 }
1393 }
1394 if self.block_cursor < self.block_data.len() {
1396 let (consumed, records_decoded) = self
1397 .decoder
1398 .decode_block(&self.block_data[self.block_cursor..], self.block_count)?;
1399 self.block_cursor += consumed;
1400 self.block_count -= records_decoded;
1401 }
1402 }
1403 self.decoder.flush_block()
1404 }
1405}
1406
1407impl<R: BufRead> Iterator for Reader<R> {
1408 type Item = Result<RecordBatch, ArrowError>;
1409
1410 fn next(&mut self) -> Option<Self::Item> {
1411 self.read().map_err(ArrowError::from).transpose()
1412 }
1413}
1414
1415impl<R: BufRead> RecordBatchReader for Reader<R> {
1416 fn schema(&self) -> SchemaRef {
1417 self.schema()
1418 }
1419}
1420
1421#[cfg(test)]
1422mod test {
1423 use crate::codec::{AvroFieldBuilder, Tz};
1424 use crate::reader::header::HeaderDecoder;
1425 use crate::reader::record::RecordDecoder;
1426 use crate::reader::{Decoder, Reader, ReaderBuilder};
1427 use crate::schema::{
1428 AVRO_ENUM_SYMBOLS_METADATA_KEY, AVRO_NAME_METADATA_KEY, AVRO_NAMESPACE_METADATA_KEY,
1429 AvroSchema, CONFLUENT_MAGIC, Fingerprint, FingerprintAlgorithm, PrimitiveType,
1430 SINGLE_OBJECT_MAGIC, SchemaStore,
1431 };
1432 use crate::test_util::arrow_test_data;
1433 use crate::writer::AvroWriter;
1434 use arrow_array::builder::{
1435 ArrayBuilder, BooleanBuilder, Float32Builder, Int32Builder, Int64Builder, ListBuilder,
1436 MapBuilder, StringBuilder, StructBuilder,
1437 };
1438 #[cfg(feature = "snappy")]
1439 use arrow_array::builder::{Float64Builder, MapFieldNames};
1440 use arrow_array::cast::AsArray;
1441 #[cfg(not(feature = "avro_custom_types"))]
1442 use arrow_array::types::Int64Type;
1443 #[cfg(feature = "avro_custom_types")]
1444 use arrow_array::types::{
1445 DurationMicrosecondType, DurationMillisecondType, DurationNanosecondType,
1446 DurationSecondType,
1447 };
1448 use arrow_array::types::{Int32Type, IntervalMonthDayNanoType};
1449 use arrow_array::*;
1450 #[cfg(feature = "snappy")]
1451 use arrow_buffer::{Buffer, NullBuffer};
1452 use arrow_buffer::{IntervalMonthDayNano, OffsetBuffer, ScalarBuffer, i256};
1453 #[cfg(feature = "avro_custom_types")]
1454 use arrow_schema::{
1455 ArrowError, DataType, Field, FieldRef, Fields, IntervalUnit, Schema, TimeUnit, UnionFields,
1456 UnionMode,
1457 };
1458 #[cfg(not(feature = "avro_custom_types"))]
1459 use arrow_schema::{
1460 ArrowError, DataType, Field, FieldRef, Fields, IntervalUnit, Schema, UnionFields, UnionMode,
1461 };
1462 use bytes::Bytes;
1463 use futures::executor::block_on;
1464 use futures::{Stream, StreamExt, TryStreamExt, stream};
1465 use serde_json::{Value, json};
1466 use std::collections::HashMap;
1467 use std::fs::File;
1468 use std::io::{BufReader, Cursor};
1469 use std::sync::Arc;
1470
1471 fn files() -> impl Iterator<Item = &'static str> {
1472 [
1473 #[cfg(feature = "snappy")]
1475 "avro/alltypes_plain.avro",
1476 #[cfg(feature = "snappy")]
1477 "avro/alltypes_plain.snappy.avro",
1478 #[cfg(feature = "zstd")]
1479 "avro/alltypes_plain.zstandard.avro",
1480 #[cfg(feature = "bzip2")]
1481 "avro/alltypes_plain.bzip2.avro",
1482 #[cfg(feature = "xz")]
1483 "avro/alltypes_plain.xz.avro",
1484 ]
1485 .into_iter()
1486 }
1487
1488 fn read_file(path: &str, batch_size: usize, utf8_view: bool) -> RecordBatch {
1489 let file = File::open(path).unwrap();
1490 let reader = ReaderBuilder::new()
1491 .with_batch_size(batch_size)
1492 .with_utf8_view(utf8_view)
1493 .build(BufReader::new(file))
1494 .unwrap();
1495 let schema = reader.schema();
1496 let batches = reader.collect::<Result<Vec<_>, _>>().unwrap();
1497 arrow::compute::concat_batches(&schema, &batches).unwrap()
1498 }
1499
1500 #[test]
1501 fn test_block_sync_marker_mismatch_errors() {
1502 let path = arrow_test_data("avro/alltypes_plain.avro");
1503 let mut bytes = std::fs::read(&path).unwrap();
1504 let last = bytes.len() - 1;
1506 bytes[last] ^= 0xFF;
1507 let reader = ReaderBuilder::new()
1508 .with_batch_size(1024)
1509 .build(std::io::Cursor::new(bytes))
1510 .unwrap();
1511 let err = reader
1512 .collect::<Result<Vec<_>, _>>()
1513 .expect_err("corrupted block sync marker should fail the read");
1514 assert!(err.to_string().contains("sync marker"), "{err}");
1515 }
1516
1517 fn read_file_strict(
1518 path: &str,
1519 batch_size: usize,
1520 utf8_view: bool,
1521 ) -> Result<Reader<BufReader<File>>, ArrowError> {
1522 let file = File::open(path)?;
1523 ReaderBuilder::new()
1524 .with_batch_size(batch_size)
1525 .with_utf8_view(utf8_view)
1526 .with_strict_mode(true)
1527 .build(BufReader::new(file))
1528 }
1529
1530 fn decode_stream<S: Stream<Item = Bytes> + Unpin>(
1531 mut decoder: Decoder,
1532 mut input: S,
1533 ) -> impl Stream<Item = Result<RecordBatch, ArrowError>> {
1534 async_stream::try_stream! {
1535 if let Some(data) = input.next().await {
1536 let consumed = decoder.decode(&data)?;
1537 if consumed < data.len() {
1538 Err(ArrowError::ParseError(
1539 "did not consume all bytes".to_string(),
1540 ))?;
1541 }
1542 }
1543 if let Some(batch) = decoder.flush()? {
1544 yield batch
1545 }
1546 }
1547 }
1548
1549 fn make_record_schema(pt: PrimitiveType) -> AvroSchema {
1550 let js = format!(
1551 r#"{{"type":"record","name":"TestRecord","fields":[{{"name":"a","type":"{}"}}]}}"#,
1552 pt.as_ref()
1553 );
1554 AvroSchema::new(js)
1555 }
1556
1557 fn make_two_schema_store() -> (
1558 SchemaStore,
1559 Fingerprint,
1560 Fingerprint,
1561 AvroSchema,
1562 AvroSchema,
1563 ) {
1564 let schema_int = make_record_schema(PrimitiveType::Int);
1565 let schema_long = make_record_schema(PrimitiveType::Long);
1566 let mut store = SchemaStore::new();
1567 let fp_int = store
1568 .register(schema_int.clone())
1569 .expect("register int schema");
1570 let fp_long = store
1571 .register(schema_long.clone())
1572 .expect("register long schema");
1573 (store, fp_int, fp_long, schema_int, schema_long)
1574 }
1575
1576 fn make_prefix(fp: Fingerprint) -> Vec<u8> {
1577 match fp {
1578 Fingerprint::Rabin(v) => {
1579 let mut out = Vec::with_capacity(2 + 8);
1580 out.extend_from_slice(&SINGLE_OBJECT_MAGIC);
1581 out.extend_from_slice(&v.to_le_bytes());
1582 out
1583 }
1584 Fingerprint::Id(v) => {
1585 panic!("make_prefix expects a Rabin fingerprint, got ({v})");
1586 }
1587 Fingerprint::Id64(v) => {
1588 panic!("make_prefix expects a Rabin fingerprint, got ({v})");
1589 }
1590 #[cfg(feature = "md5")]
1591 Fingerprint::MD5(v) => {
1592 panic!("make_prefix expects a Rabin fingerprint, got ({v:?})");
1593 }
1594 #[cfg(feature = "sha256")]
1595 Fingerprint::SHA256(id) => {
1596 panic!("make_prefix expects a Rabin fingerprint, got ({id:?})");
1597 }
1598 }
1599 }
1600
1601 fn make_decoder(store: &SchemaStore, fp: Fingerprint, reader_schema: &AvroSchema) -> Decoder {
1602 ReaderBuilder::new()
1603 .with_batch_size(8)
1604 .with_reader_schema(reader_schema.clone())
1605 .with_writer_schema_store(store.clone())
1606 .with_active_fingerprint(fp)
1607 .build_decoder()
1608 .expect("decoder")
1609 }
1610
1611 fn make_id_prefix(id: u32, additional: usize) -> Vec<u8> {
1612 let capacity = CONFLUENT_MAGIC.len() + size_of::<u32>() + additional;
1613 let mut out = Vec::with_capacity(capacity);
1614 out.extend_from_slice(&CONFLUENT_MAGIC);
1615 out.extend_from_slice(&id.to_be_bytes());
1616 out
1617 }
1618
1619 fn make_message_id(id: u32, value: i64) -> Vec<u8> {
1620 let encoded_value = encode_zigzag(value);
1621 let mut msg = make_id_prefix(id, encoded_value.len());
1622 msg.extend_from_slice(&encoded_value);
1623 msg
1624 }
1625
1626 fn make_id64_prefix(id: u64, additional: usize) -> Vec<u8> {
1627 let capacity = CONFLUENT_MAGIC.len() + size_of::<u64>() + additional;
1628 let mut out = Vec::with_capacity(capacity);
1629 out.extend_from_slice(&CONFLUENT_MAGIC);
1630 out.extend_from_slice(&id.to_be_bytes());
1631 out
1632 }
1633
1634 fn make_message_id64(id: u64, value: i64) -> Vec<u8> {
1635 let encoded_value = encode_zigzag(value);
1636 let mut msg = make_id64_prefix(id, encoded_value.len());
1637 msg.extend_from_slice(&encoded_value);
1638 msg
1639 }
1640
1641 fn make_value_schema(pt: PrimitiveType) -> AvroSchema {
1642 let json_schema = format!(
1643 r#"{{"type":"record","name":"S","fields":[{{"name":"v","type":"{}"}}]}}"#,
1644 pt.as_ref()
1645 );
1646 AvroSchema::new(json_schema)
1647 }
1648
1649 fn encode_zigzag(value: i64) -> Vec<u8> {
1650 let mut n = ((value << 1) ^ (value >> 63)) as u64;
1651 let mut out = Vec::new();
1652 loop {
1653 if (n & !0x7F) == 0 {
1654 out.push(n as u8);
1655 break;
1656 } else {
1657 out.push(((n & 0x7F) | 0x80) as u8);
1658 n >>= 7;
1659 }
1660 }
1661 out
1662 }
1663
1664 fn make_message(fp: Fingerprint, value: i64) -> Vec<u8> {
1665 let mut msg = make_prefix(fp);
1666 msg.extend_from_slice(&encode_zigzag(value));
1667 msg
1668 }
1669
1670 fn load_writer_schema_json(path: &str) -> Value {
1671 let file = File::open(path).unwrap();
1672 let (header, _) = super::read_header(BufReader::new(file)).unwrap();
1673 let schema = header.schema().unwrap().unwrap();
1674 serde_json::to_value(&schema).unwrap()
1675 }
1676
1677 fn make_reader_schema_with_promotions(
1678 path: &str,
1679 promotions: &HashMap<&str, &str>,
1680 ) -> AvroSchema {
1681 let mut root = load_writer_schema_json(path);
1682 assert_eq!(root["type"], "record", "writer schema must be a record");
1683 let fields = root
1684 .get_mut("fields")
1685 .and_then(|f| f.as_array_mut())
1686 .expect("record has fields");
1687 for f in fields.iter_mut() {
1688 let Some(name) = f.get("name").and_then(|n| n.as_str()) else {
1689 continue;
1690 };
1691 if let Some(new_ty) = promotions.get(name) {
1692 let ty = f.get_mut("type").expect("field has a type");
1693 match ty {
1694 Value::String(_) => {
1695 *ty = Value::String((*new_ty).to_string());
1696 }
1697 Value::Array(arr) => {
1699 for b in arr.iter_mut() {
1700 match b {
1701 Value::String(s) if s != "null" => {
1702 *b = Value::String((*new_ty).to_string());
1703 break;
1704 }
1705 Value::Object(_) => {
1706 *b = Value::String((*new_ty).to_string());
1707 break;
1708 }
1709 _ => {}
1710 }
1711 }
1712 }
1713 Value::Object(_) => {
1714 *ty = Value::String((*new_ty).to_string());
1715 }
1716 _ => {}
1717 }
1718 }
1719 }
1720 AvroSchema::new(root.to_string())
1721 }
1722
1723 fn make_reader_schema_with_enum_remap(
1724 path: &str,
1725 remap: &HashMap<&str, Vec<&str>>,
1726 ) -> AvroSchema {
1727 let mut root = load_writer_schema_json(path);
1728 assert_eq!(root["type"], "record", "writer schema must be a record");
1729 let fields = root
1730 .get_mut("fields")
1731 .and_then(|f| f.as_array_mut())
1732 .expect("record has fields");
1733
1734 fn to_symbols_array(symbols: &[&str]) -> Value {
1735 Value::Array(symbols.iter().map(|s| Value::String((*s).into())).collect())
1736 }
1737
1738 fn update_enum_symbols(ty: &mut Value, symbols: &Value) {
1739 match ty {
1740 Value::Object(map) => {
1741 if matches!(map.get("type"), Some(Value::String(t)) if t == "enum") {
1742 map.insert("symbols".to_string(), symbols.clone());
1743 }
1744 }
1745 Value::Array(arr) => {
1746 for b in arr.iter_mut() {
1747 if let Value::Object(map) = b
1748 && matches!(map.get("type"), Some(Value::String(t)) if t == "enum")
1749 {
1750 map.insert("symbols".to_string(), symbols.clone());
1751 }
1752 }
1753 }
1754 _ => {}
1755 }
1756 }
1757 for f in fields.iter_mut() {
1758 let Some(name) = f.get("name").and_then(|n| n.as_str()) else {
1759 continue;
1760 };
1761 if let Some(new_symbols) = remap.get(name) {
1762 let symbols_val = to_symbols_array(new_symbols);
1763 let ty = f.get_mut("type").expect("field has a type");
1764 update_enum_symbols(ty, &symbols_val);
1765 }
1766 }
1767 AvroSchema::new(root.to_string())
1768 }
1769
1770 fn read_alltypes_with_reader_schema(path: &str, reader_schema: AvroSchema) -> RecordBatch {
1771 let file = File::open(path).unwrap();
1772 let reader = ReaderBuilder::new()
1773 .with_batch_size(1024)
1774 .with_utf8_view(false)
1775 .with_reader_schema(reader_schema)
1776 .build(BufReader::new(file))
1777 .unwrap();
1778 let schema = reader.schema();
1779 let batches = reader.collect::<Result<Vec<_>, _>>().unwrap();
1780 arrow::compute::concat_batches(&schema, &batches).unwrap()
1781 }
1782
1783 fn make_reader_schema_with_selected_fields_in_order(
1784 path: &str,
1785 selected: &[&str],
1786 ) -> AvroSchema {
1787 let mut root = load_writer_schema_json(path);
1788 assert_eq!(root["type"], "record", "writer schema must be a record");
1789 let writer_fields = root
1790 .get("fields")
1791 .and_then(|f| f.as_array())
1792 .expect("record has fields");
1793 let mut field_map: HashMap<String, Value> = HashMap::with_capacity(writer_fields.len());
1794 for f in writer_fields {
1795 if let Some(name) = f.get("name").and_then(|n| n.as_str()) {
1796 field_map.insert(name.to_string(), f.clone());
1797 }
1798 }
1799 let mut new_fields = Vec::with_capacity(selected.len());
1800 for name in selected {
1801 let f = field_map
1802 .get(*name)
1803 .unwrap_or_else(|| panic!("field '{name}' not found in writer schema"))
1804 .clone();
1805 new_fields.push(f);
1806 }
1807 root["fields"] = Value::Array(new_fields);
1808 AvroSchema::new(root.to_string())
1809 }
1810
1811 fn write_ocf(schema: &Schema, batches: &[RecordBatch]) -> Vec<u8> {
1812 let mut w = AvroWriter::new(Vec::<u8>::new(), schema.clone()).expect("writer");
1813 for b in batches {
1814 w.write(b).expect("write");
1815 }
1816 w.finish().expect("finish");
1817 w.into_inner()
1818 }
1819
1820 #[test]
1821 fn ocf_projection_no_reader_schema_reorder() -> Result<(), Box<dyn std::error::Error>> {
1822 let writer_schema = Schema::new(vec![
1824 Field::new("id", DataType::Int32, false),
1825 Field::new("name", DataType::Utf8, false),
1826 Field::new("is_active", DataType::Boolean, false),
1827 ]);
1828 let batch = RecordBatch::try_new(
1829 Arc::new(writer_schema.clone()),
1830 vec![
1831 Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef,
1832 Arc::new(StringArray::from(vec!["a", "b"])) as ArrayRef,
1833 Arc::new(BooleanArray::from(vec![true, false])) as ArrayRef,
1834 ],
1835 )?;
1836 let bytes = write_ocf(&writer_schema, &[batch]);
1837 let mut reader = ReaderBuilder::new()
1839 .with_projection(vec![2, 0])
1840 .build(Cursor::new(bytes))?;
1841 let out = reader.next().unwrap()?;
1842 assert_eq!(out.num_columns(), 2);
1843 assert_eq!(out.schema().field(0).name(), "is_active");
1844 assert_eq!(out.schema().field(1).name(), "id");
1845 let is_active = out.column(0).as_boolean();
1846 assert!(is_active.value(0));
1847 assert!(!is_active.value(1));
1848 let id = out.column(1).as_primitive::<Int32Type>();
1849 assert_eq!(id.value(0), 1);
1850 assert_eq!(id.value(1), 2);
1851 Ok(())
1852 }
1853
1854 #[test]
1855 fn ocf_projection_with_reader_schema_alias_and_default()
1856 -> Result<(), Box<dyn std::error::Error>> {
1857 let writer_schema = Schema::new(vec![
1859 Field::new("id", DataType::Int64, false),
1860 Field::new("name", DataType::Utf8, false),
1861 ]);
1862 let batch = RecordBatch::try_new(
1863 Arc::new(writer_schema.clone()),
1864 vec![
1865 Arc::new(Int64Array::from(vec![1, 2])) as ArrayRef,
1866 Arc::new(StringArray::from(vec!["a", "b"])) as ArrayRef,
1867 ],
1868 )?;
1869 let bytes = write_ocf(&writer_schema, &[batch]);
1870 let reader_json = r#"
1874 {
1875 "type": "record",
1876 "name": "topLevelRecord",
1877 "fields": [
1878 { "name": "id", "type": "long" },
1879 { "name": "full_name", "type": ["null","string"], "aliases": ["name"], "default": null },
1880 { "name": "is_active", "type": "boolean", "default": true }
1881 ]
1882 }"#;
1883 let mut reader = ReaderBuilder::new()
1885 .with_reader_schema(AvroSchema::new(reader_json.to_string()))
1886 .with_projection(vec![1, 2])
1887 .build(Cursor::new(bytes))?;
1888 let out = reader.next().unwrap()?;
1889 assert_eq!(out.num_columns(), 2);
1890 assert_eq!(out.schema().field(0).name(), "full_name");
1891 assert_eq!(out.schema().field(1).name(), "is_active");
1892 let full_name = out.column(0).as_string::<i32>();
1893 assert_eq!(full_name.value(0), "a");
1894 assert_eq!(full_name.value(1), "b");
1895 let is_active = out.column(1).as_boolean();
1896 assert!(is_active.value(0));
1897 assert!(is_active.value(1));
1898 Ok(())
1899 }
1900
1901 #[test]
1902 fn projection_errors_out_of_bounds_and_duplicate() -> Result<(), Box<dyn std::error::Error>> {
1903 let writer_schema = Schema::new(vec![
1904 Field::new("a", DataType::Int32, false),
1905 Field::new("b", DataType::Int32, false),
1906 ]);
1907 let batch = RecordBatch::try_new(
1908 Arc::new(writer_schema.clone()),
1909 vec![
1910 Arc::new(Int32Array::from(vec![1])) as ArrayRef,
1911 Arc::new(Int32Array::from(vec![2])) as ArrayRef,
1912 ],
1913 )?;
1914 let bytes = write_ocf(&writer_schema, &[batch]);
1915 let err = ReaderBuilder::new()
1916 .with_projection(vec![2])
1917 .build(Cursor::new(bytes.clone()))
1918 .unwrap_err();
1919 assert!(matches!(err, ArrowError::AvroError(_)));
1920 assert!(err.to_string().contains("out of bounds"));
1921 let err = ReaderBuilder::new()
1922 .with_projection(vec![0, 0])
1923 .build(Cursor::new(bytes))
1924 .unwrap_err();
1925 assert!(matches!(err, ArrowError::AvroError(_)));
1926 assert!(err.to_string().contains("Duplicate projection index"));
1927 Ok(())
1928 }
1929
1930 #[test]
1931 #[cfg(feature = "snappy")]
1932 fn test_alltypes_plain_with_projection_and_reader_schema() {
1933 use std::fs::File;
1934 use std::io::BufReader;
1935 let path = arrow_test_data("avro/alltypes_plain.avro");
1936 let reader_schema = make_reader_schema_with_selected_fields_in_order(
1938 &path,
1939 &["double_col", "id", "tinyint_col"],
1940 );
1941 let file = File::open(&path).expect("open avro/alltypes_plain.avro");
1942 let reader = ReaderBuilder::new()
1943 .with_batch_size(1024)
1944 .with_reader_schema(reader_schema)
1945 .with_projection(vec![1, 2]) .build(BufReader::new(file))
1947 .expect("build reader with projection and reader schema");
1948 let schema = reader.schema();
1949 assert_eq!(schema.fields().len(), 2);
1951 assert_eq!(schema.field(0).name(), "id");
1952 assert_eq!(schema.field(1).name(), "tinyint_col");
1953 let batches: Vec<RecordBatch> = reader.collect::<Result<Vec<_>, _>>().unwrap();
1954 assert_eq!(batches.len(), 1);
1955 let batch = &batches[0];
1956 assert_eq!(batch.num_rows(), 8);
1957 assert_eq!(batch.num_columns(), 2);
1958 let expected = RecordBatch::try_from_iter_with_nullable([
1962 (
1963 "id",
1964 Arc::new(Int32Array::from(vec![4, 5, 6, 7, 2, 3, 0, 1])) as ArrayRef,
1965 true,
1966 ),
1967 (
1968 "tinyint_col",
1969 Arc::new(Int32Array::from(vec![0, 1, 0, 1, 0, 1, 0, 1])) as ArrayRef,
1970 true,
1971 ),
1972 ])
1973 .unwrap();
1974 assert_eq!(
1975 batch, &expected,
1976 "Projected batch mismatch for alltypes_plain.avro with reader schema and projection [1, 2]"
1977 );
1978 }
1979
1980 #[test]
1981 #[cfg(feature = "snappy")]
1982 fn test_alltypes_plain_with_projection() {
1983 use std::fs::File;
1984 use std::io::BufReader;
1985 let path = arrow_test_data("avro/alltypes_plain.avro");
1986 let file = File::open(&path).expect("open avro/alltypes_plain.avro");
1987 let reader = ReaderBuilder::new()
1988 .with_batch_size(1024)
1989 .with_projection(vec![2, 0, 5])
1990 .build(BufReader::new(file))
1991 .expect("build reader with projection");
1992 let schema = reader.schema();
1993 assert_eq!(schema.fields().len(), 3);
1994 assert_eq!(schema.field(0).name(), "tinyint_col");
1995 assert_eq!(schema.field(1).name(), "id");
1996 assert_eq!(schema.field(2).name(), "bigint_col");
1997 let batches: Vec<RecordBatch> = reader.collect::<Result<Vec<_>, _>>().unwrap();
1998 assert_eq!(batches.len(), 1);
1999 let batch = &batches[0];
2000 assert_eq!(batch.num_rows(), 8);
2001 assert_eq!(batch.num_columns(), 3);
2002 let expected = RecordBatch::try_from_iter_with_nullable([
2003 (
2004 "tinyint_col",
2005 Arc::new(Int32Array::from(vec![0, 1, 0, 1, 0, 1, 0, 1])) as ArrayRef,
2006 true,
2007 ),
2008 (
2009 "id",
2010 Arc::new(Int32Array::from(vec![4, 5, 6, 7, 2, 3, 0, 1])) as ArrayRef,
2011 true,
2012 ),
2013 (
2014 "bigint_col",
2015 Arc::new(Int64Array::from(vec![0, 10, 0, 10, 0, 10, 0, 10])) as ArrayRef,
2016 true,
2017 ),
2018 ])
2019 .unwrap();
2020 assert_eq!(
2021 batch, &expected,
2022 "Projected batch mismatch for alltypes_plain.avro with projection [2, 0, 5]"
2023 );
2024 }
2025
2026 #[test]
2027 fn writer_string_reader_nullable_with_alias() -> Result<(), Box<dyn std::error::Error>> {
2028 let writer_schema = Schema::new(vec![
2029 Field::new("id", DataType::Int64, false),
2030 Field::new("name", DataType::Utf8, false),
2031 ]);
2032 let batch = RecordBatch::try_new(
2033 Arc::new(writer_schema.clone()),
2034 vec![
2035 Arc::new(Int64Array::from(vec![1, 2])) as ArrayRef,
2036 Arc::new(StringArray::from(vec!["a", "b"])) as ArrayRef,
2037 ],
2038 )?;
2039 let bytes = write_ocf(&writer_schema, &[batch]);
2040 let reader_json = r#"
2041 {
2042 "type": "record",
2043 "name": "topLevelRecord",
2044 "fields": [
2045 { "name": "id", "type": "long" },
2046 { "name": "full_name", "type": ["null","string"], "aliases": ["name"], "default": null },
2047 { "name": "is_active", "type": "boolean", "default": true }
2048 ]
2049 }"#;
2050 let mut reader = ReaderBuilder::new()
2051 .with_reader_schema(AvroSchema::new(reader_json.to_string()))
2052 .build(Cursor::new(bytes))?;
2053 let out = reader.next().unwrap()?;
2054 let full_name = out.column(1).as_string::<i32>();
2055 assert_eq!(full_name.value(0), "a");
2056 assert_eq!(full_name.value(1), "b");
2057 Ok(())
2058 }
2059
2060 #[test]
2061 fn writer_string_reader_string_null_order_second() -> Result<(), Box<dyn std::error::Error>> {
2062 let writer_schema = Schema::new(vec![Field::new("name", DataType::Utf8, false)]);
2064 let batch = RecordBatch::try_new(
2065 Arc::new(writer_schema.clone()),
2066 vec![Arc::new(StringArray::from(vec!["x", "y"])) as ArrayRef],
2067 )?;
2068 let bytes = write_ocf(&writer_schema, &[batch]);
2069
2070 let reader_json = r#"
2072 {
2073 "type":"record", "name":"topLevelRecord",
2074 "fields":[ { "name":"name", "type":["string","null"], "default":"x" } ]
2075 }"#;
2076
2077 let mut reader = ReaderBuilder::new()
2078 .with_reader_schema(AvroSchema::new(reader_json.to_string()))
2079 .build(Cursor::new(bytes))?;
2080
2081 let out = reader.next().unwrap()?;
2082 assert_eq!(out.num_rows(), 2);
2083
2084 let name = out.column(0).as_string::<i32>();
2086 assert_eq!(name.value(0), "x");
2087 assert_eq!(name.value(1), "y");
2088
2089 Ok(())
2090 }
2091
2092 #[test]
2093 fn promotion_writer_int_reader_nullable_long() -> Result<(), Box<dyn std::error::Error>> {
2094 let writer_schema = Schema::new(vec![Field::new("v", DataType::Int32, false)]);
2096 let batch = RecordBatch::try_new(
2097 Arc::new(writer_schema.clone()),
2098 vec![Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef],
2099 )?;
2100 let bytes = write_ocf(&writer_schema, &[batch]);
2101
2102 let reader_json = r#"
2104 {
2105 "type":"record", "name":"topLevelRecord",
2106 "fields":[ { "name":"v", "type":["null","long"], "default": null } ]
2107 }"#;
2108
2109 let mut reader = ReaderBuilder::new()
2110 .with_reader_schema(AvroSchema::new(reader_json.to_string()))
2111 .build(Cursor::new(bytes))?;
2112
2113 let out = reader.next().unwrap()?;
2114 assert_eq!(out.num_rows(), 3);
2115
2116 let v = out
2118 .column(0)
2119 .as_primitive::<arrow_array::types::Int64Type>();
2120 assert_eq!(v.values(), &[1, 2, 3]);
2121 assert!(
2122 out.column(0).nulls().is_none(),
2123 "expected no validity bitmap for all-valid column"
2124 );
2125
2126 Ok(())
2127 }
2128
2129 #[test]
2130 fn test_alltypes_schema_promotion_mixed() {
2131 for file in files() {
2132 let file = arrow_test_data(file);
2133 let mut promotions: HashMap<&str, &str> = HashMap::new();
2134 promotions.insert("id", "long");
2135 promotions.insert("tinyint_col", "float");
2136 promotions.insert("smallint_col", "double");
2137 promotions.insert("int_col", "double");
2138 promotions.insert("bigint_col", "double");
2139 promotions.insert("float_col", "double");
2140 promotions.insert("date_string_col", "string");
2141 promotions.insert("string_col", "string");
2142 let reader_schema = make_reader_schema_with_promotions(&file, &promotions);
2143 let batch = read_alltypes_with_reader_schema(&file, reader_schema);
2144 let expected = RecordBatch::try_from_iter_with_nullable([
2145 (
2146 "id",
2147 Arc::new(Int64Array::from(vec![4i64, 5, 6, 7, 2, 3, 0, 1])) as _,
2148 true,
2149 ),
2150 (
2151 "bool_col",
2152 Arc::new(BooleanArray::from_iter((0..8).map(|x| Some(x % 2 == 0)))) as _,
2153 true,
2154 ),
2155 (
2156 "tinyint_col",
2157 Arc::new(Float32Array::from_iter_values(
2158 (0..8).map(|x| (x % 2) as f32),
2159 )) as _,
2160 true,
2161 ),
2162 (
2163 "smallint_col",
2164 Arc::new(Float64Array::from_iter_values(
2165 (0..8).map(|x| (x % 2) as f64),
2166 )) as _,
2167 true,
2168 ),
2169 (
2170 "int_col",
2171 Arc::new(Float64Array::from_iter_values(
2172 (0..8).map(|x| (x % 2) as f64),
2173 )) as _,
2174 true,
2175 ),
2176 (
2177 "bigint_col",
2178 Arc::new(Float64Array::from_iter_values(
2179 (0..8).map(|x| ((x % 2) * 10) as f64),
2180 )) as _,
2181 true,
2182 ),
2183 (
2184 "float_col",
2185 Arc::new(Float64Array::from_iter_values(
2186 (0..8).map(|x| ((x % 2) as f32 * 1.1f32) as f64),
2187 )) as _,
2188 true,
2189 ),
2190 (
2191 "double_col",
2192 Arc::new(Float64Array::from_iter_values(
2193 (0..8).map(|x| (x % 2) as f64 * 10.1),
2194 )) as _,
2195 true,
2196 ),
2197 (
2198 "date_string_col",
2199 Arc::new(StringArray::from(vec![
2200 "03/01/09", "03/01/09", "04/01/09", "04/01/09", "02/01/09", "02/01/09",
2201 "01/01/09", "01/01/09",
2202 ])) as _,
2203 true,
2204 ),
2205 (
2206 "string_col",
2207 Arc::new(StringArray::from(
2208 (0..8)
2209 .map(|x| if x % 2 == 0 { "0" } else { "1" })
2210 .collect::<Vec<_>>(),
2211 )) as _,
2212 true,
2213 ),
2214 (
2215 "timestamp_col",
2216 Arc::new(
2217 TimestampMicrosecondArray::from_iter_values([
2218 1235865600000000, 1235865660000000, 1238544000000000, 1238544060000000, 1233446400000000, 1233446460000000, 1230768000000000, 1230768060000000, ])
2227 .with_timezone("+00:00"),
2228 ) as _,
2229 true,
2230 ),
2231 ])
2232 .unwrap();
2233 assert_eq!(batch, expected, "mismatch for file {file}");
2234 }
2235 }
2236
2237 #[test]
2238 fn test_alltypes_schema_promotion_long_to_float_only() {
2239 for file in files() {
2240 let file = arrow_test_data(file);
2241 let mut promotions: HashMap<&str, &str> = HashMap::new();
2242 promotions.insert("bigint_col", "float");
2243 let reader_schema = make_reader_schema_with_promotions(&file, &promotions);
2244 let batch = read_alltypes_with_reader_schema(&file, reader_schema);
2245 let expected = RecordBatch::try_from_iter_with_nullable([
2246 (
2247 "id",
2248 Arc::new(Int32Array::from(vec![4, 5, 6, 7, 2, 3, 0, 1])) as _,
2249 true,
2250 ),
2251 (
2252 "bool_col",
2253 Arc::new(BooleanArray::from_iter((0..8).map(|x| Some(x % 2 == 0)))) as _,
2254 true,
2255 ),
2256 (
2257 "tinyint_col",
2258 Arc::new(Int32Array::from_iter_values((0..8).map(|x| x % 2))) as _,
2259 true,
2260 ),
2261 (
2262 "smallint_col",
2263 Arc::new(Int32Array::from_iter_values((0..8).map(|x| x % 2))) as _,
2264 true,
2265 ),
2266 (
2267 "int_col",
2268 Arc::new(Int32Array::from_iter_values((0..8).map(|x| x % 2))) as _,
2269 true,
2270 ),
2271 (
2272 "bigint_col",
2273 Arc::new(Float32Array::from_iter_values(
2274 (0..8).map(|x| ((x % 2) * 10) as f32),
2275 )) as _,
2276 true,
2277 ),
2278 (
2279 "float_col",
2280 Arc::new(Float32Array::from_iter_values(
2281 (0..8).map(|x| (x % 2) as f32 * 1.1),
2282 )) as _,
2283 true,
2284 ),
2285 (
2286 "double_col",
2287 Arc::new(Float64Array::from_iter_values(
2288 (0..8).map(|x| (x % 2) as f64 * 10.1),
2289 )) as _,
2290 true,
2291 ),
2292 (
2293 "date_string_col",
2294 Arc::new(BinaryArray::from_iter_values([
2295 [48, 51, 47, 48, 49, 47, 48, 57],
2296 [48, 51, 47, 48, 49, 47, 48, 57],
2297 [48, 52, 47, 48, 49, 47, 48, 57],
2298 [48, 52, 47, 48, 49, 47, 48, 57],
2299 [48, 50, 47, 48, 49, 47, 48, 57],
2300 [48, 50, 47, 48, 49, 47, 48, 57],
2301 [48, 49, 47, 48, 49, 47, 48, 57],
2302 [48, 49, 47, 48, 49, 47, 48, 57],
2303 ])) as _,
2304 true,
2305 ),
2306 (
2307 "string_col",
2308 Arc::new(BinaryArray::from_iter_values((0..8).map(|x| [48 + x % 2]))) as _,
2309 true,
2310 ),
2311 (
2312 "timestamp_col",
2313 Arc::new(
2314 TimestampMicrosecondArray::from_iter_values([
2315 1235865600000000, 1235865660000000, 1238544000000000, 1238544060000000, 1233446400000000, 1233446460000000, 1230768000000000, 1230768060000000, ])
2324 .with_timezone("+00:00"),
2325 ) as _,
2326 true,
2327 ),
2328 ])
2329 .unwrap();
2330 assert_eq!(batch, expected, "mismatch for file {file}");
2331 }
2332 }
2333
2334 #[test]
2335 fn test_alltypes_schema_promotion_bytes_to_string_only() {
2336 for file in files() {
2337 let file = arrow_test_data(file);
2338 let mut promotions: HashMap<&str, &str> = HashMap::new();
2339 promotions.insert("date_string_col", "string");
2340 promotions.insert("string_col", "string");
2341 let reader_schema = make_reader_schema_with_promotions(&file, &promotions);
2342 let batch = read_alltypes_with_reader_schema(&file, reader_schema);
2343 let expected = RecordBatch::try_from_iter_with_nullable([
2344 (
2345 "id",
2346 Arc::new(Int32Array::from(vec![4, 5, 6, 7, 2, 3, 0, 1])) as _,
2347 true,
2348 ),
2349 (
2350 "bool_col",
2351 Arc::new(BooleanArray::from_iter((0..8).map(|x| Some(x % 2 == 0)))) as _,
2352 true,
2353 ),
2354 (
2355 "tinyint_col",
2356 Arc::new(Int32Array::from_iter_values((0..8).map(|x| x % 2))) as _,
2357 true,
2358 ),
2359 (
2360 "smallint_col",
2361 Arc::new(Int32Array::from_iter_values((0..8).map(|x| x % 2))) as _,
2362 true,
2363 ),
2364 (
2365 "int_col",
2366 Arc::new(Int32Array::from_iter_values((0..8).map(|x| x % 2))) as _,
2367 true,
2368 ),
2369 (
2370 "bigint_col",
2371 Arc::new(Int64Array::from_iter_values((0..8).map(|x| (x % 2) * 10))) as _,
2372 true,
2373 ),
2374 (
2375 "float_col",
2376 Arc::new(Float32Array::from_iter_values(
2377 (0..8).map(|x| (x % 2) as f32 * 1.1),
2378 )) as _,
2379 true,
2380 ),
2381 (
2382 "double_col",
2383 Arc::new(Float64Array::from_iter_values(
2384 (0..8).map(|x| (x % 2) as f64 * 10.1),
2385 )) as _,
2386 true,
2387 ),
2388 (
2389 "date_string_col",
2390 Arc::new(StringArray::from(vec![
2391 "03/01/09", "03/01/09", "04/01/09", "04/01/09", "02/01/09", "02/01/09",
2392 "01/01/09", "01/01/09",
2393 ])) as _,
2394 true,
2395 ),
2396 (
2397 "string_col",
2398 Arc::new(StringArray::from(
2399 (0..8)
2400 .map(|x| if x % 2 == 0 { "0" } else { "1" })
2401 .collect::<Vec<_>>(),
2402 )) as _,
2403 true,
2404 ),
2405 (
2406 "timestamp_col",
2407 Arc::new(
2408 TimestampMicrosecondArray::from_iter_values([
2409 1235865600000000, 1235865660000000, 1238544000000000, 1238544060000000, 1233446400000000, 1233446460000000, 1230768000000000, 1230768060000000, ])
2418 .with_timezone("+00:00"),
2419 ) as _,
2420 true,
2421 ),
2422 ])
2423 .unwrap();
2424 assert_eq!(batch, expected, "mismatch for file {file}");
2425 }
2426 }
2427
2428 #[test]
2429 #[cfg(feature = "snappy")]
2431 fn test_alltypes_illegal_promotion_bool_to_double_errors() {
2432 let file = arrow_test_data("avro/alltypes_plain.avro");
2433 let mut promotions: HashMap<&str, &str> = HashMap::new();
2434 promotions.insert("bool_col", "double"); let reader_schema = make_reader_schema_with_promotions(&file, &promotions);
2436 let file_handle = File::open(&file).unwrap();
2437 let result = ReaderBuilder::new()
2438 .with_reader_schema(reader_schema)
2439 .build(BufReader::new(file_handle));
2440 let err = result.expect_err("expected illegal promotion to error");
2441 let msg = err.to_string();
2442 assert!(
2443 msg.contains("Illegal promotion") || msg.contains("illegal promotion"),
2444 "unexpected error: {msg}"
2445 );
2446 }
2447
2448 #[test]
2449 fn test_simple_enum_with_reader_schema_mapping() {
2450 let file = arrow_test_data("avro/simple_enum.avro");
2451 let mut remap: HashMap<&str, Vec<&str>> = HashMap::new();
2452 remap.insert("f1", vec!["d", "c", "b", "a"]);
2453 remap.insert("f2", vec!["h", "g", "f", "e"]);
2454 remap.insert("f3", vec!["k", "i", "j"]);
2455 let reader_schema = make_reader_schema_with_enum_remap(&file, &remap);
2456 let actual = read_alltypes_with_reader_schema(&file, reader_schema);
2457 let dict_type = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
2458 let f1_keys = Int32Array::from(vec![3, 2, 1, 0]);
2460 let f1_vals = StringArray::from(vec!["d", "c", "b", "a"]);
2461 let f1 = DictionaryArray::<Int32Type>::try_new(f1_keys, Arc::new(f1_vals)).unwrap();
2462 let mut md_f1 = HashMap::new();
2463 md_f1.insert(
2464 AVRO_ENUM_SYMBOLS_METADATA_KEY.to_string(),
2465 r#"["d","c","b","a"]"#.to_string(),
2466 );
2467 md_f1.insert("avro.name".to_string(), "enum1".to_string());
2469 md_f1.insert("avro.namespace".to_string(), "ns1".to_string());
2470 let f1_field = Field::new("f1", dict_type.clone(), false).with_metadata(md_f1);
2471 let f2_keys = Int32Array::from(vec![1, 0, 3, 2]);
2473 let f2_vals = StringArray::from(vec!["h", "g", "f", "e"]);
2474 let f2 = DictionaryArray::<Int32Type>::try_new(f2_keys, Arc::new(f2_vals)).unwrap();
2475 let mut md_f2 = HashMap::new();
2476 md_f2.insert(
2477 AVRO_ENUM_SYMBOLS_METADATA_KEY.to_string(),
2478 r#"["h","g","f","e"]"#.to_string(),
2479 );
2480 md_f2.insert("avro.name".to_string(), "enum2".to_string());
2482 md_f2.insert("avro.namespace".to_string(), "ns2".to_string());
2483 let f2_field = Field::new("f2", dict_type.clone(), false).with_metadata(md_f2);
2484 let f3_keys = Int32Array::from(vec![Some(2), Some(0), None, Some(1)]);
2486 let f3_vals = StringArray::from(vec!["k", "i", "j"]);
2487 let f3 = DictionaryArray::<Int32Type>::try_new(f3_keys, Arc::new(f3_vals)).unwrap();
2488 let mut md_f3 = HashMap::new();
2489 md_f3.insert(
2490 AVRO_ENUM_SYMBOLS_METADATA_KEY.to_string(),
2491 r#"["k","i","j"]"#.to_string(),
2492 );
2493 md_f3.insert("avro.name".to_string(), "enum3".to_string());
2495 md_f3.insert("avro.namespace".to_string(), "ns1".to_string());
2496 let f3_field = Field::new("f3", dict_type.clone(), true).with_metadata(md_f3);
2497 let expected_schema = Arc::new(Schema::new(vec![f1_field, f2_field, f3_field]));
2498 let expected = RecordBatch::try_new(
2499 expected_schema,
2500 vec![Arc::new(f1) as ArrayRef, Arc::new(f2), Arc::new(f3)],
2501 )
2502 .unwrap();
2503 assert_eq!(actual, expected);
2504 }
2505
2506 #[test]
2507 fn test_schema_store_register_lookup() {
2508 let schema_int = make_record_schema(PrimitiveType::Int);
2509 let schema_long = make_record_schema(PrimitiveType::Long);
2510 let mut store = SchemaStore::new();
2511 let fp_int = store.register(schema_int.clone()).unwrap();
2512 let fp_long = store.register(schema_long.clone()).unwrap();
2513 assert_eq!(store.lookup(&fp_int).cloned(), Some(schema_int));
2514 assert_eq!(store.lookup(&fp_long).cloned(), Some(schema_long));
2515 assert_eq!(store.fingerprint_algorithm(), FingerprintAlgorithm::Rabin);
2516 }
2517
2518 #[test]
2519 fn test_unknown_fingerprint_is_error() {
2520 let (store, fp_int, _fp_long, _schema_int, schema_long) = make_two_schema_store();
2521 let unknown_fp = Fingerprint::Rabin(0xDEAD_BEEF_DEAD_BEEF);
2522 let prefix = make_prefix(unknown_fp);
2523 let mut decoder = make_decoder(&store, fp_int, &schema_long);
2524 let err = decoder.decode(&prefix).expect_err("decode should error");
2525 let msg = err.to_string();
2526 assert!(
2527 msg.contains("Unknown fingerprint"),
2528 "unexpected message: {msg}"
2529 );
2530 }
2531
2532 #[test]
2533 fn test_handle_prefix_incomplete_magic() {
2534 let (store, fp_int, _fp_long, _schema_int, schema_long) = make_two_schema_store();
2535 let mut decoder = make_decoder(&store, fp_int, &schema_long);
2536 let buf = &SINGLE_OBJECT_MAGIC[..1];
2537 let res = decoder.handle_prefix(buf).unwrap();
2538 assert_eq!(res, Some(0));
2539 assert!(decoder.pending_schema.is_none());
2540 }
2541
2542 #[test]
2543 fn test_handle_prefix_magic_mismatch() {
2544 let (store, fp_int, _fp_long, _schema_int, schema_long) = make_two_schema_store();
2545 let mut decoder = make_decoder(&store, fp_int, &schema_long);
2546 let buf = [0xFFu8, 0x00u8, 0x01u8];
2547 let res = decoder.handle_prefix(&buf).unwrap();
2548 assert!(res.is_none());
2549 }
2550
2551 #[test]
2552 fn test_handle_prefix_incomplete_fingerprint() {
2553 let (store, fp_int, fp_long, _schema_int, schema_long) = make_two_schema_store();
2554 let mut decoder = make_decoder(&store, fp_int, &schema_long);
2555 let long_bytes = match fp_long {
2556 Fingerprint::Rabin(v) => v.to_le_bytes(),
2557 Fingerprint::Id(id) => panic!("expected Rabin fingerprint, got ({id})"),
2558 Fingerprint::Id64(id) => panic!("expected Rabin fingerprint, got ({id})"),
2559 #[cfg(feature = "md5")]
2560 Fingerprint::MD5(v) => panic!("expected Rabin fingerprint, got ({v:?})"),
2561 #[cfg(feature = "sha256")]
2562 Fingerprint::SHA256(v) => panic!("expected Rabin fingerprint, got ({v:?})"),
2563 };
2564 let mut buf = Vec::from(SINGLE_OBJECT_MAGIC);
2565 buf.extend_from_slice(&long_bytes[..4]);
2566 let res = decoder.handle_prefix(&buf).unwrap();
2567 assert_eq!(res, Some(0));
2568 assert!(decoder.pending_schema.is_none());
2569 }
2570
2571 #[test]
2572 fn test_handle_prefix_valid_prefix_switches_schema() {
2573 let (store, fp_int, fp_long, _schema_int, schema_long) = make_two_schema_store();
2574 let mut decoder = make_decoder(&store, fp_int, &schema_long);
2575 let writer_schema_long = schema_long.schema().unwrap();
2576 let root_long = AvroFieldBuilder::new(&writer_schema_long).build().unwrap();
2577 let long_decoder = RecordDecoder::try_new_with_options(root_long.data_type()).unwrap();
2578 let _ = decoder.cache.insert(fp_long, long_decoder);
2579 let mut buf = Vec::from(SINGLE_OBJECT_MAGIC);
2580 match fp_long {
2581 Fingerprint::Rabin(v) => buf.extend_from_slice(&v.to_le_bytes()),
2582 Fingerprint::Id(id) => panic!("expected Rabin fingerprint, got ({id})"),
2583 Fingerprint::Id64(id) => panic!("expected Rabin fingerprint, got ({id})"),
2584 #[cfg(feature = "md5")]
2585 Fingerprint::MD5(v) => panic!("expected Rabin fingerprint, got ({v:?})"),
2586 #[cfg(feature = "sha256")]
2587 Fingerprint::SHA256(v) => panic!("expected Rabin fingerprint, got ({v:?})"),
2588 }
2589 let consumed = decoder.handle_prefix(&buf).unwrap().unwrap();
2590 assert_eq!(consumed, buf.len());
2591 assert!(decoder.pending_schema.is_some());
2592 assert_eq!(decoder.pending_schema.as_ref().unwrap().0, fp_long);
2593 }
2594
2595 #[test]
2596 fn test_decoder_projection_multiple_writer_schemas_no_reader_schema()
2597 -> Result<(), Box<dyn std::error::Error>> {
2598 let writer_v1 = AvroSchema::new(
2600 r#"{"type":"record","name":"E","fields":[{"name":"a","type":"int"},{"name":"b","type":"string"}]}"#
2601 .to_string(),
2602 );
2603 let writer_v2 = AvroSchema::new(
2604 r#"{"type":"record","name":"E","fields":[{"name":"a","type":"long"},{"name":"b","type":"string"},{"name":"c","type":"int"}]}"#
2605 .to_string(),
2606 );
2607 let mut store = SchemaStore::new();
2608 let fp1 = store.register(writer_v1)?;
2609 let fp2 = store.register(writer_v2)?;
2610 let mut decoder = ReaderBuilder::new()
2611 .with_writer_schema_store(store)
2612 .with_active_fingerprint(fp1)
2613 .with_batch_size(8)
2614 .with_projection(vec![1])
2615 .build_decoder()?;
2616 let mut msg1 = make_prefix(fp1);
2618 msg1.extend_from_slice(&encode_zigzag(1)); msg1.push((1u8) << 1);
2620 msg1.extend_from_slice(b"x");
2621 let mut msg2 = make_prefix(fp2);
2623 msg2.extend_from_slice(&encode_zigzag(2)); msg2.push((1u8) << 1);
2625 msg2.extend_from_slice(b"y");
2626 msg2.extend_from_slice(&encode_zigzag(7)); decoder.decode(&msg1)?;
2628 let batch1 = decoder.flush()?.expect("batch1");
2629 assert_eq!(batch1.num_columns(), 1);
2630 assert_eq!(batch1.schema().field(0).name(), "b");
2631 let b1 = batch1.column(0).as_string::<i32>();
2632 assert_eq!(b1.value(0), "x");
2633 decoder.decode(&msg2)?;
2634 let batch2 = decoder.flush()?.expect("batch2");
2635 assert_eq!(batch2.num_columns(), 1);
2636 assert_eq!(batch2.schema().field(0).name(), "b");
2637 let b2 = batch2.column(0).as_string::<i32>();
2638 assert_eq!(b2.value(0), "y");
2639 Ok(())
2640 }
2641
2642 #[test]
2643 fn test_two_messages_same_schema() {
2644 let writer_schema = make_value_schema(PrimitiveType::Int);
2645 let reader_schema = writer_schema.clone();
2646 let mut store = SchemaStore::new();
2647 let fp = store.register(writer_schema).unwrap();
2648 let msg1 = make_message(fp, 42);
2649 let msg2 = make_message(fp, 11);
2650 let input = [msg1.clone(), msg2.clone()].concat();
2651 let mut decoder = ReaderBuilder::new()
2652 .with_batch_size(8)
2653 .with_reader_schema(reader_schema.clone())
2654 .with_writer_schema_store(store)
2655 .with_active_fingerprint(fp)
2656 .build_decoder()
2657 .unwrap();
2658 let _ = decoder.decode(&input).unwrap();
2659 let batch = decoder.flush().unwrap().expect("batch");
2660 assert_eq!(batch.num_rows(), 2);
2661 let col = batch
2662 .column(0)
2663 .as_any()
2664 .downcast_ref::<Int32Array>()
2665 .unwrap();
2666 assert_eq!(col.value(0), 42);
2667 assert_eq!(col.value(1), 11);
2668 }
2669
2670 #[test]
2671 fn test_two_messages_schema_switch() {
2672 let w_int = make_value_schema(PrimitiveType::Int);
2673 let w_long = make_value_schema(PrimitiveType::Long);
2674 let mut store = SchemaStore::new();
2675 let fp_int = store.register(w_int).unwrap();
2676 let fp_long = store.register(w_long).unwrap();
2677 let msg_int = make_message(fp_int, 1);
2678 let msg_long = make_message(fp_long, 123456789_i64);
2679 let mut decoder = ReaderBuilder::new()
2680 .with_batch_size(8)
2681 .with_writer_schema_store(store)
2682 .with_active_fingerprint(fp_int)
2683 .build_decoder()
2684 .unwrap();
2685 let _ = decoder.decode(&msg_int).unwrap();
2686 let batch1 = decoder.flush().unwrap().expect("batch1");
2687 assert_eq!(batch1.num_rows(), 1);
2688 assert_eq!(
2689 batch1
2690 .column(0)
2691 .as_any()
2692 .downcast_ref::<Int32Array>()
2693 .unwrap()
2694 .value(0),
2695 1
2696 );
2697 let _ = decoder.decode(&msg_long).unwrap();
2698 let batch2 = decoder.flush().unwrap().expect("batch2");
2699 assert_eq!(batch2.num_rows(), 1);
2700 assert_eq!(
2701 batch2
2702 .column(0)
2703 .as_any()
2704 .downcast_ref::<Int64Array>()
2705 .unwrap()
2706 .value(0),
2707 123456789_i64
2708 );
2709 }
2710
2711 #[test]
2712 fn test_two_messages_same_schema_id() {
2713 let writer_schema = make_value_schema(PrimitiveType::Int);
2714 let reader_schema = writer_schema.clone();
2715 let id = 100u32;
2716 let mut store = SchemaStore::new_with_type(FingerprintAlgorithm::Id);
2718 let _ = store
2719 .set(Fingerprint::Id(id), writer_schema.clone())
2720 .expect("set id schema");
2721 let msg1 = make_message_id(id, 21);
2722 let msg2 = make_message_id(id, 22);
2723 let input = [msg1.clone(), msg2.clone()].concat();
2724 let mut decoder = ReaderBuilder::new()
2725 .with_batch_size(8)
2726 .with_reader_schema(reader_schema)
2727 .with_writer_schema_store(store)
2728 .with_active_fingerprint(Fingerprint::Id(id))
2729 .build_decoder()
2730 .unwrap();
2731 let _ = decoder.decode(&input).unwrap();
2732 let batch = decoder.flush().unwrap().expect("batch");
2733 assert_eq!(batch.num_rows(), 2);
2734 let col = batch
2735 .column(0)
2736 .as_any()
2737 .downcast_ref::<Int32Array>()
2738 .unwrap();
2739 assert_eq!(col.value(0), 21);
2740 assert_eq!(col.value(1), 22);
2741 }
2742
2743 #[test]
2744 fn test_unknown_id_fingerprint_is_error() {
2745 let writer_schema = make_value_schema(PrimitiveType::Int);
2746 let id_known = 7u32;
2747 let id_unknown = 9u32;
2748 let mut store = SchemaStore::new_with_type(FingerprintAlgorithm::Id);
2749 let _ = store
2750 .set(Fingerprint::Id(id_known), writer_schema.clone())
2751 .expect("set id schema");
2752 let mut decoder = ReaderBuilder::new()
2753 .with_batch_size(8)
2754 .with_reader_schema(writer_schema)
2755 .with_writer_schema_store(store)
2756 .with_active_fingerprint(Fingerprint::Id(id_known))
2757 .build_decoder()
2758 .unwrap();
2759 let prefix = make_id_prefix(id_unknown, 0);
2760 let err = decoder.decode(&prefix).expect_err("decode should error");
2761 let msg = err.to_string();
2762 assert!(
2763 msg.contains("Unknown fingerprint"),
2764 "unexpected message: {msg}"
2765 );
2766 }
2767
2768 #[test]
2769 fn test_handle_prefix_id_incomplete_magic() {
2770 let writer_schema = make_value_schema(PrimitiveType::Int);
2771 let id = 5u32;
2772 let mut store = SchemaStore::new_with_type(FingerprintAlgorithm::Id);
2773 let _ = store
2774 .set(Fingerprint::Id(id), writer_schema.clone())
2775 .expect("set id schema");
2776 let mut decoder = ReaderBuilder::new()
2777 .with_batch_size(8)
2778 .with_reader_schema(writer_schema)
2779 .with_writer_schema_store(store)
2780 .with_active_fingerprint(Fingerprint::Id(id))
2781 .build_decoder()
2782 .unwrap();
2783 let buf = &CONFLUENT_MAGIC[..0]; let res = decoder.handle_prefix(buf).unwrap();
2785 assert_eq!(res, Some(0));
2786 assert!(decoder.pending_schema.is_none());
2787 }
2788
2789 #[test]
2790 fn test_two_messages_same_schema_id64() {
2791 let writer_schema = make_value_schema(PrimitiveType::Int);
2792 let reader_schema = writer_schema.clone();
2793 let id = 100u64;
2794 let mut store = SchemaStore::new_with_type(FingerprintAlgorithm::Id64);
2796 let _ = store
2797 .set(Fingerprint::Id64(id), writer_schema.clone())
2798 .expect("set id schema");
2799 let msg1 = make_message_id64(id, 21);
2800 let msg2 = make_message_id64(id, 22);
2801 let input = [msg1.clone(), msg2.clone()].concat();
2802 let mut decoder = ReaderBuilder::new()
2803 .with_batch_size(8)
2804 .with_reader_schema(reader_schema)
2805 .with_writer_schema_store(store)
2806 .with_active_fingerprint(Fingerprint::Id64(id))
2807 .build_decoder()
2808 .unwrap();
2809 let _ = decoder.decode(&input).unwrap();
2810 let batch = decoder.flush().unwrap().expect("batch");
2811 assert_eq!(batch.num_rows(), 2);
2812 let col = batch
2813 .column(0)
2814 .as_any()
2815 .downcast_ref::<Int32Array>()
2816 .unwrap();
2817 assert_eq!(col.value(0), 21);
2818 assert_eq!(col.value(1), 22);
2819 }
2820
2821 #[test]
2822 fn test_decode_stream_with_schema() {
2823 struct TestCase<'a> {
2824 name: &'a str,
2825 schema: &'a str,
2826 expected_error: Option<&'a str>,
2827 }
2828 let tests = vec![
2829 TestCase {
2830 name: "success",
2831 schema: r#"{"type":"record","name":"test","fields":[{"name":"f2","type":"string"}]}"#,
2832 expected_error: None,
2833 },
2834 TestCase {
2835 name: "valid schema invalid data",
2836 schema: r#"{"type":"record","name":"test","fields":[{"name":"f2","type":"long"}]}"#,
2837 expected_error: Some("did not consume all bytes"),
2838 },
2839 ];
2840 for test in tests {
2841 let avro_schema = AvroSchema::new(test.schema.to_string());
2842 let mut store = SchemaStore::new();
2843 let fp = store.register(avro_schema.clone()).unwrap();
2844 let prefix = make_prefix(fp);
2845 let record_val = "some_string";
2846 let mut body = prefix;
2847 body.push((record_val.len() as u8) << 1);
2848 body.extend_from_slice(record_val.as_bytes());
2849 let decoder_res = ReaderBuilder::new()
2850 .with_batch_size(1)
2851 .with_writer_schema_store(store)
2852 .with_active_fingerprint(fp)
2853 .build_decoder();
2854 let decoder = match decoder_res {
2855 Ok(d) => d,
2856 Err(e) => {
2857 if let Some(expected) = test.expected_error {
2858 assert!(
2859 e.to_string().contains(expected),
2860 "Test '{}' failed at build – expected '{expected}', got '{e}'",
2861 test.name
2862 );
2863 continue;
2864 } else {
2865 panic!("Test '{}' failed during build: {e}", test.name);
2866 }
2867 }
2868 };
2869 let stream = Box::pin(stream::once(async { Bytes::from(body) }));
2870 let decoded_stream = decode_stream(decoder, stream);
2871 let batches_result: Result<Vec<RecordBatch>, ArrowError> =
2872 block_on(decoded_stream.try_collect());
2873 match (batches_result, test.expected_error) {
2874 (Ok(batches), None) => {
2875 let batch =
2876 arrow::compute::concat_batches(&batches[0].schema(), &batches).unwrap();
2877 let expected_field = Field::new("f2", DataType::Utf8, false);
2878 let expected_schema = Arc::new(Schema::new(vec![expected_field]));
2879 let expected_array = Arc::new(StringArray::from(vec![record_val]));
2880 let expected_batch =
2881 RecordBatch::try_new(expected_schema, vec![expected_array]).unwrap();
2882 assert_eq!(batch, expected_batch, "Test '{}'", test.name);
2883 }
2884 (Err(e), Some(expected)) => {
2885 assert!(
2886 e.to_string().contains(expected),
2887 "Test '{}' – expected error containing '{expected}', got '{e}'",
2888 test.name
2889 );
2890 }
2891 (Ok(_), Some(expected)) => {
2892 panic!(
2893 "Test '{}' expected failure ('{expected}') but succeeded",
2894 test.name
2895 );
2896 }
2897 (Err(e), None) => {
2898 panic!("Test '{}' unexpectedly failed with '{e}'", test.name);
2899 }
2900 }
2901 }
2902 }
2903
2904 #[test]
2905 fn test_utf8view_support() {
2906 struct TestHelper;
2907 impl TestHelper {
2908 fn with_utf8view(field: &Field) -> Field {
2909 match field.data_type() {
2910 DataType::Utf8 => {
2911 Field::new(field.name(), DataType::Utf8View, field.is_nullable())
2912 .with_metadata(field.metadata().clone())
2913 }
2914 _ => field.clone(),
2915 }
2916 }
2917 }
2918
2919 let field = TestHelper::with_utf8view(&Field::new("str_field", DataType::Utf8, false));
2920
2921 assert_eq!(field.data_type(), &DataType::Utf8View);
2922
2923 let array = StringViewArray::from(vec!["test1", "test2"]);
2924 let batch =
2925 RecordBatch::try_from_iter(vec![("str_field", Arc::new(array) as ArrayRef)]).unwrap();
2926
2927 assert!(batch.column(0).as_any().is::<StringViewArray>());
2928 }
2929
2930 fn make_reader_schema_with_default_fields(
2931 path: &str,
2932 default_fields: Vec<Value>,
2933 ) -> AvroSchema {
2934 let mut root = load_writer_schema_json(path);
2935 assert_eq!(root["type"], "record", "writer schema must be a record");
2936 root.as_object_mut()
2937 .expect("schema is a JSON object")
2938 .insert("fields".to_string(), Value::Array(default_fields));
2939 AvroSchema::new(root.to_string())
2940 }
2941
2942 #[test]
2943 fn test_schema_resolution_defaults_all_supported_types() {
2944 let path = "test/data/skippable_types.avro";
2945 let duration_default = "\u{0000}".repeat(12);
2946 let reader_schema = make_reader_schema_with_default_fields(
2947 path,
2948 vec![
2949 serde_json::json!({"name":"d_bool","type":"boolean","default":true}),
2950 serde_json::json!({"name":"d_int","type":"int","default":42}),
2951 serde_json::json!({"name":"d_long","type":"long","default":12345}),
2952 serde_json::json!({"name":"d_float","type":"float","default":1.5}),
2953 serde_json::json!({"name":"d_double","type":"double","default":2.25}),
2954 serde_json::json!({"name":"d_bytes","type":"bytes","default":"XYZ"}),
2955 serde_json::json!({"name":"d_string","type":"string","default":"hello"}),
2956 serde_json::json!({"name":"d_date","type":{"type":"int","logicalType":"date"},"default":0}),
2957 serde_json::json!({"name":"d_time_ms","type":{"type":"int","logicalType":"time-millis"},"default":1000}),
2958 serde_json::json!({"name":"d_time_us","type":{"type":"long","logicalType":"time-micros"},"default":2000}),
2959 serde_json::json!({"name":"d_ts_ms","type":{"type":"long","logicalType":"local-timestamp-millis"},"default":0}),
2960 serde_json::json!({"name":"d_ts_us","type":{"type":"long","logicalType":"local-timestamp-micros"},"default":0}),
2961 serde_json::json!({"name":"d_decimal","type":{"type":"bytes","logicalType":"decimal","precision":10,"scale":2},"default":""}),
2962 serde_json::json!({"name":"d_fixed","type":{"type":"fixed","name":"F4","size":4},"default":"ABCD"}),
2963 serde_json::json!({"name":"d_enum","type":{"type":"enum","name":"E","symbols":["A","B","C"]},"default":"A"}),
2964 serde_json::json!({"name":"d_duration","type":{"type":"fixed","name":"Dur","size":12,"logicalType":"duration"},"default":duration_default}),
2965 serde_json::json!({"name":"d_uuid","type":{"type":"string","logicalType":"uuid"},"default":"00000000-0000-0000-0000-000000000000"}),
2966 serde_json::json!({"name":"d_array","type":{"type":"array","items":"int"},"default":[1,2,3]}),
2967 serde_json::json!({"name":"d_map","type":{"type":"map","values":"long"},"default":{"a":1,"b":2}}),
2968 serde_json::json!({"name":"d_record","type":{
2969 "type":"record","name":"DefaultRec","fields":[
2970 {"name":"x","type":"int"},
2971 {"name":"y","type":["null","string"],"default":null}
2972 ]
2973 },"default":{"x":7}}),
2974 serde_json::json!({"name":"d_nullable_null","type":["null","int"],"default":null}),
2975 serde_json::json!({"name":"d_nullable_value","type":["int","null"],"default":123}),
2976 ],
2977 );
2978 let actual = read_alltypes_with_reader_schema(path, reader_schema);
2979 let num_rows = actual.num_rows();
2980 assert!(num_rows > 0, "skippable_types.avro should contain rows");
2981 assert_eq!(
2982 actual.num_columns(),
2983 22,
2984 "expected exactly our defaulted fields"
2985 );
2986 let mut arrays: Vec<Arc<dyn Array>> = Vec::with_capacity(22);
2987 arrays.push(Arc::new(BooleanArray::from_iter(std::iter::repeat_n(
2988 Some(true),
2989 num_rows,
2990 ))));
2991 arrays.push(Arc::new(Int32Array::from_iter_values(std::iter::repeat_n(
2992 42, num_rows,
2993 ))));
2994 arrays.push(Arc::new(Int64Array::from_iter_values(std::iter::repeat_n(
2995 12345, num_rows,
2996 ))));
2997 arrays.push(Arc::new(Float32Array::from_iter_values(
2998 std::iter::repeat_n(1.5f32, num_rows),
2999 )));
3000 arrays.push(Arc::new(Float64Array::from_iter_values(
3001 std::iter::repeat_n(2.25f64, num_rows),
3002 )));
3003 arrays.push(Arc::new(BinaryArray::from_iter_values(
3004 std::iter::repeat_n(b"XYZ".as_ref(), num_rows),
3005 )));
3006 arrays.push(Arc::new(StringArray::from_iter_values(
3007 std::iter::repeat_n("hello", num_rows),
3008 )));
3009 arrays.push(Arc::new(Date32Array::from_iter_values(
3010 std::iter::repeat_n(0, num_rows),
3011 )));
3012 arrays.push(Arc::new(Time32MillisecondArray::from_iter_values(
3013 std::iter::repeat_n(1_000, num_rows),
3014 )));
3015 arrays.push(Arc::new(Time64MicrosecondArray::from_iter_values(
3016 std::iter::repeat_n(2_000i64, num_rows),
3017 )));
3018 arrays.push(Arc::new(TimestampMillisecondArray::from_iter_values(
3019 std::iter::repeat_n(0i64, num_rows),
3020 )));
3021 arrays.push(Arc::new(TimestampMicrosecondArray::from_iter_values(
3022 std::iter::repeat_n(0i64, num_rows),
3023 )));
3024 #[cfg(feature = "small_decimals")]
3025 let decimal = Decimal64Array::from_iter_values(std::iter::repeat_n(0i64, num_rows))
3026 .with_precision_and_scale(10, 2)
3027 .unwrap();
3028 #[cfg(not(feature = "small_decimals"))]
3029 let decimal = Decimal128Array::from_iter_values(std::iter::repeat_n(0i128, num_rows))
3030 .with_precision_and_scale(10, 2)
3031 .unwrap();
3032 arrays.push(Arc::new(decimal));
3033 let fixed_iter = std::iter::repeat_n(Some(*b"ABCD"), num_rows);
3034 arrays.push(Arc::new(
3035 FixedSizeBinaryArray::try_from_sparse_iter_with_size(fixed_iter, 4).unwrap(),
3036 ));
3037 let enum_keys = Int32Array::from_iter_values(std::iter::repeat_n(0, num_rows));
3038 let enum_values = StringArray::from_iter_values(["A", "B", "C"]);
3039 let enum_arr =
3040 DictionaryArray::<Int32Type>::try_new(enum_keys, Arc::new(enum_values)).unwrap();
3041 arrays.push(Arc::new(enum_arr));
3042 let duration_values = std::iter::repeat_n(
3043 Some(IntervalMonthDayNanoType::make_value(0, 0, 0)),
3044 num_rows,
3045 );
3046 let duration_arr: IntervalMonthDayNanoArray = duration_values.collect();
3047 arrays.push(Arc::new(duration_arr));
3048 let uuid_bytes = [0u8; 16];
3049 let uuid_iter = std::iter::repeat_n(Some(uuid_bytes), num_rows);
3050 arrays.push(Arc::new(
3051 FixedSizeBinaryArray::try_from_sparse_iter_with_size(uuid_iter, 16).unwrap(),
3052 ));
3053 let item_field = Arc::new(Field::new(
3054 Field::LIST_FIELD_DEFAULT_NAME,
3055 DataType::Int32,
3056 false,
3057 ));
3058 let mut list_builder = ListBuilder::new(Int32Builder::new()).with_field(item_field);
3059 for _ in 0..num_rows {
3060 list_builder.values().append_value(1);
3061 list_builder.values().append_value(2);
3062 list_builder.values().append_value(3);
3063 list_builder.append(true);
3064 }
3065 arrays.push(Arc::new(list_builder.finish()));
3066 let values_field = Arc::new(Field::new(
3067 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
3068 DataType::Int64,
3069 false,
3070 ));
3071 let mut map_builder = MapBuilder::new(
3072 Some(builder::MapFieldNames {
3073 entry: Field::MAP_ENTRIES_FIELD_DEFAULT_NAME.to_string(),
3074 key: Field::MAP_KEY_FIELD_DEFAULT_NAME.to_string(),
3075 value: Field::MAP_VALUE_FIELD_DEFAULT_NAME.to_string(),
3076 }),
3077 StringBuilder::new(),
3078 Int64Builder::new(),
3079 )
3080 .with_values_field(values_field);
3081 for _ in 0..num_rows {
3082 let (keys, vals) = map_builder.entries();
3083 keys.append_value("a");
3084 vals.append_value(1);
3085 keys.append_value("b");
3086 vals.append_value(2);
3087 map_builder.append(true).unwrap();
3088 }
3089 arrays.push(Arc::new(map_builder.finish()));
3090 let rec_fields: Fields = Fields::from(vec![
3091 Field::new("x", DataType::Int32, false),
3092 Field::new("y", DataType::Utf8, true),
3093 ]);
3094 let mut sb = StructBuilder::new(
3095 rec_fields.clone(),
3096 vec![
3097 Box::new(Int32Builder::new()),
3098 Box::new(StringBuilder::new()),
3099 ],
3100 );
3101 for _ in 0..num_rows {
3102 sb.field_builder::<Int32Builder>(0).unwrap().append_value(7);
3103 sb.field_builder::<StringBuilder>(1).unwrap().append_null();
3104 sb.append(true);
3105 }
3106 arrays.push(Arc::new(sb.finish()));
3107 arrays.push(Arc::new(Int32Array::from_iter(std::iter::repeat_n(
3108 None::<i32>,
3109 num_rows,
3110 ))));
3111 arrays.push(Arc::new(Int32Array::from_iter_values(std::iter::repeat_n(
3112 123, num_rows,
3113 ))));
3114 let expected = RecordBatch::try_new(actual.schema(), arrays).unwrap();
3115 assert_eq!(
3116 actual, expected,
3117 "defaults should materialize correctly for all fields"
3118 );
3119 }
3120
3121 #[test]
3122 fn test_schema_resolution_default_enum_invalid_symbol_errors() {
3123 let path = "test/data/skippable_types.avro";
3124 let bad_schema = make_reader_schema_with_default_fields(
3125 path,
3126 vec![serde_json::json!({
3127 "name":"bad_enum",
3128 "type":{"type":"enum","name":"E","symbols":["A","B","C"]},
3129 "default":"Z"
3130 })],
3131 );
3132 let file = File::open(path).unwrap();
3133 let res = ReaderBuilder::new()
3134 .with_reader_schema(bad_schema)
3135 .build(BufReader::new(file));
3136 let err = res.expect_err("expected enum default validation to fail");
3137 let msg = err.to_string();
3138 let lower_msg = msg.to_lowercase();
3139 assert!(
3140 lower_msg.contains("enum")
3141 && (lower_msg.contains("symbol") || lower_msg.contains("default")),
3142 "unexpected error: {msg}"
3143 );
3144 }
3145
3146 #[test]
3147 fn test_schema_resolution_default_fixed_size_mismatch_errors() {
3148 let path = "test/data/skippable_types.avro";
3149 let bad_schema = make_reader_schema_with_default_fields(
3150 path,
3151 vec![serde_json::json!({
3152 "name":"bad_fixed",
3153 "type":{"type":"fixed","name":"F","size":4},
3154 "default":"ABC"
3155 })],
3156 );
3157 let file = File::open(path).unwrap();
3158 let res = ReaderBuilder::new()
3159 .with_reader_schema(bad_schema)
3160 .build(BufReader::new(file));
3161 let err = res.expect_err("expected fixed default validation to fail");
3162 let msg = err.to_string();
3163 let lower_msg = msg.to_lowercase();
3164 assert!(
3165 lower_msg.contains("fixed")
3166 && (lower_msg.contains("size")
3167 || lower_msg.contains("length")
3168 || lower_msg.contains("does not match")),
3169 "unexpected error: {msg}"
3170 );
3171 }
3172
3173 #[test]
3174 fn test_timestamp_with_utc_tz() {
3175 let path = arrow_test_data("avro/alltypes_plain.avro");
3176 let reader_schema =
3177 make_reader_schema_with_selected_fields_in_order(&path, &["timestamp_col"]);
3178 let file = File::open(path).unwrap();
3179 let reader = ReaderBuilder::new()
3180 .with_batch_size(1024)
3181 .with_utf8_view(false)
3182 .with_reader_schema(reader_schema)
3183 .with_tz(Tz::Utc)
3184 .build(BufReader::new(file))
3185 .unwrap();
3186 let schema = reader.schema();
3187 let batches = reader.collect::<Result<Vec<_>, _>>().unwrap();
3188 let batch = arrow::compute::concat_batches(&schema, &batches).unwrap();
3189 let expected = RecordBatch::try_from_iter_with_nullable([(
3190 "timestamp_col",
3191 Arc::new(
3192 TimestampMicrosecondArray::from_iter_values([
3193 1235865600000000, 1235865660000000, 1238544000000000, 1238544060000000, 1233446400000000, 1233446460000000, 1230768000000000, 1230768060000000, ])
3202 .with_timezone("UTC"),
3203 ) as _,
3204 true,
3205 )])
3206 .unwrap();
3207 assert_eq!(batch, expected);
3208 }
3209
3210 #[test]
3211 #[cfg(feature = "snappy")]
3213 fn test_alltypes_skip_writer_fields_keep_double_only() {
3214 let file = arrow_test_data("avro/alltypes_plain.avro");
3215 let reader_schema =
3216 make_reader_schema_with_selected_fields_in_order(&file, &["double_col"]);
3217 let batch = read_alltypes_with_reader_schema(&file, reader_schema);
3218 let expected = RecordBatch::try_from_iter_with_nullable([(
3219 "double_col",
3220 Arc::new(Float64Array::from_iter_values(
3221 (0..8).map(|x| (x % 2) as f64 * 10.1),
3222 )) as _,
3223 true,
3224 )])
3225 .unwrap();
3226 assert_eq!(batch, expected);
3227 }
3228
3229 #[test]
3230 #[cfg(feature = "snappy")]
3232 fn test_alltypes_skip_writer_fields_reorder_and_skip_many() {
3233 let file = arrow_test_data("avro/alltypes_plain.avro");
3234 let reader_schema =
3235 make_reader_schema_with_selected_fields_in_order(&file, &["timestamp_col", "id"]);
3236 let batch = read_alltypes_with_reader_schema(&file, reader_schema);
3237 let expected = RecordBatch::try_from_iter_with_nullable([
3238 (
3239 "timestamp_col",
3240 Arc::new(
3241 TimestampMicrosecondArray::from_iter_values([
3242 1235865600000000, 1235865660000000, 1238544000000000, 1238544060000000, 1233446400000000, 1233446460000000, 1230768000000000, 1230768060000000, ])
3251 .with_timezone("+00:00"),
3252 ) as _,
3253 true,
3254 ),
3255 (
3256 "id",
3257 Arc::new(Int32Array::from(vec![4, 5, 6, 7, 2, 3, 0, 1])) as _,
3258 true,
3259 ),
3260 ])
3261 .unwrap();
3262 assert_eq!(batch, expected);
3263 }
3264
3265 #[test]
3266 fn test_skippable_types_project_each_field_individually() {
3267 let path = "test/data/skippable_types.avro";
3268 let full = read_file(path, 1024, false);
3269 let schema_full = full.schema();
3270 let num_rows = full.num_rows();
3271 let writer_json = load_writer_schema_json(path);
3272 assert_eq!(
3273 writer_json["type"], "record",
3274 "writer schema must be a record"
3275 );
3276 let fields_json = writer_json
3277 .get("fields")
3278 .and_then(|f| f.as_array())
3279 .expect("record has fields");
3280 assert_eq!(
3281 schema_full.fields().len(),
3282 fields_json.len(),
3283 "full read column count vs writer fields"
3284 );
3285 fn rebuild_list_array_with_element(
3286 col: &ArrayRef,
3287 new_elem: Arc<Field>,
3288 is_large: bool,
3289 ) -> ArrayRef {
3290 if is_large {
3291 let list = col
3292 .as_any()
3293 .downcast_ref::<LargeListArray>()
3294 .expect("expected LargeListArray");
3295 let offsets = list.offsets().clone();
3296 let values = list.values().clone();
3297 let validity = list.nulls().cloned();
3298 Arc::new(LargeListArray::try_new(new_elem, offsets, values, validity).unwrap())
3299 } else {
3300 let list = col
3301 .as_any()
3302 .downcast_ref::<ListArray>()
3303 .expect("expected ListArray");
3304 let offsets = list.offsets().clone();
3305 let values = list.values().clone();
3306 let validity = list.nulls().cloned();
3307 Arc::new(ListArray::try_new(new_elem, offsets, values, validity).unwrap())
3308 }
3309 }
3310 for (idx, f) in fields_json.iter().enumerate() {
3311 let name = f
3312 .get("name")
3313 .and_then(|n| n.as_str())
3314 .unwrap_or_else(|| panic!("field at index {idx} has no name"));
3315 let reader_schema = make_reader_schema_with_selected_fields_in_order(path, &[name]);
3316 let projected = read_alltypes_with_reader_schema(path, reader_schema);
3317 assert_eq!(
3318 projected.num_columns(),
3319 1,
3320 "projected batch should contain exactly the selected column '{name}'"
3321 );
3322 assert_eq!(
3323 projected.num_rows(),
3324 num_rows,
3325 "row count mismatch for projected column '{name}'"
3326 );
3327 let col_full = full.column(idx).clone();
3328 let full_field = schema_full.field(idx).as_ref().clone();
3329 let proj_field_ref = projected.schema().field(0).clone();
3330 let proj_field = proj_field_ref.as_ref();
3331 let top_meta = proj_field.metadata().clone();
3332 let (expected_field_ref, expected_col): (Arc<Field>, ArrayRef) =
3333 match (full_field.data_type(), proj_field.data_type()) {
3334 (&DataType::List(_), DataType::List(proj_elem)) => {
3335 let new_col =
3336 rebuild_list_array_with_element(&col_full, proj_elem.clone(), false);
3337 let nf = Field::new(
3338 full_field.name().clone(),
3339 proj_field.data_type().clone(),
3340 full_field.is_nullable(),
3341 )
3342 .with_metadata(top_meta);
3343 (Arc::new(nf), new_col)
3344 }
3345 (&DataType::LargeList(_), DataType::LargeList(proj_elem)) => {
3346 let new_col =
3347 rebuild_list_array_with_element(&col_full, proj_elem.clone(), true);
3348 let nf = Field::new(
3349 full_field.name().clone(),
3350 proj_field.data_type().clone(),
3351 full_field.is_nullable(),
3352 )
3353 .with_metadata(top_meta);
3354 (Arc::new(nf), new_col)
3355 }
3356 _ => {
3357 let nf = full_field.with_metadata(top_meta);
3358 (Arc::new(nf), col_full)
3359 }
3360 };
3361
3362 let expected = RecordBatch::try_new(
3363 Arc::new(Schema::new(vec![expected_field_ref])),
3364 vec![expected_col],
3365 )
3366 .unwrap();
3367 assert_eq!(
3368 projected, expected,
3369 "projected column '{name}' mismatch vs full read column"
3370 );
3371 }
3372 }
3373
3374 #[test]
3375 fn test_union_fields_avro_nullable_and_general_unions() {
3376 let path = "test/data/union_fields.avro";
3377 let batch = read_file(path, 1024, false);
3378 let schema = batch.schema();
3379 let idx = schema.index_of("nullable_int_nullfirst").unwrap();
3380 let a = batch.column(idx).as_primitive::<Int32Type>();
3381 assert_eq!(a.len(), 4);
3382 assert!(a.is_null(0));
3383 assert_eq!(a.value(1), 42);
3384 assert!(a.is_null(2));
3385 assert_eq!(a.value(3), 0);
3386 let idx = schema.index_of("nullable_string_nullsecond").unwrap();
3387 let s = batch
3388 .column(idx)
3389 .as_any()
3390 .downcast_ref::<StringArray>()
3391 .expect("nullable_string_nullsecond should be Utf8");
3392 assert_eq!(s.len(), 4);
3393 assert_eq!(s.value(0), "s1");
3394 assert!(s.is_null(1));
3395 assert_eq!(s.value(2), "s3");
3396 assert!(s.is_valid(3)); assert_eq!(s.value(3), "");
3398 let idx = schema.index_of("union_prim").unwrap();
3399 let u = batch
3400 .column(idx)
3401 .as_any()
3402 .downcast_ref::<UnionArray>()
3403 .expect("union_prim should be Union");
3404 let fields = match u.data_type() {
3405 DataType::Union(fields, mode) => {
3406 assert!(matches!(mode, UnionMode::Dense), "expect dense unions");
3407 fields
3408 }
3409 other => panic!("expected Union, got {other:?}"),
3410 };
3411 let tid_by_name = |name: &str| -> i8 {
3412 for (tid, f) in fields.iter() {
3413 if f.name() == name {
3414 return tid;
3415 }
3416 }
3417 panic!("union child '{name}' not found");
3418 };
3419 let expected_type_ids = vec![
3420 tid_by_name("long"),
3421 tid_by_name("int"),
3422 tid_by_name("float"),
3423 tid_by_name("double"),
3424 ];
3425 let type_ids: Vec<i8> = u.type_ids().iter().copied().collect();
3426 assert_eq!(
3427 type_ids, expected_type_ids,
3428 "branch selection for union_prim rows"
3429 );
3430 let longs = u
3431 .child(tid_by_name("long"))
3432 .as_any()
3433 .downcast_ref::<Int64Array>()
3434 .unwrap();
3435 assert_eq!(longs.len(), 1);
3436 let ints = u
3437 .child(tid_by_name("int"))
3438 .as_any()
3439 .downcast_ref::<Int32Array>()
3440 .unwrap();
3441 assert_eq!(ints.len(), 1);
3442 let floats = u
3443 .child(tid_by_name("float"))
3444 .as_any()
3445 .downcast_ref::<Float32Array>()
3446 .unwrap();
3447 assert_eq!(floats.len(), 1);
3448 let doubles = u
3449 .child(tid_by_name("double"))
3450 .as_any()
3451 .downcast_ref::<Float64Array>()
3452 .unwrap();
3453 assert_eq!(doubles.len(), 1);
3454 let idx = schema.index_of("union_bytes_vs_string").unwrap();
3455 let u = batch
3456 .column(idx)
3457 .as_any()
3458 .downcast_ref::<UnionArray>()
3459 .expect("union_bytes_vs_string should be Union");
3460 let fields = match u.data_type() {
3461 DataType::Union(fields, _) => fields,
3462 other => panic!("expected Union, got {other:?}"),
3463 };
3464 let tid_by_name = |name: &str| -> i8 {
3465 for (tid, f) in fields.iter() {
3466 if f.name() == name {
3467 return tid;
3468 }
3469 }
3470 panic!("union child '{name}' not found");
3471 };
3472 let tid_bytes = tid_by_name("bytes");
3473 let tid_string = tid_by_name("string");
3474 let type_ids: Vec<i8> = u.type_ids().iter().copied().collect();
3475 assert_eq!(
3476 type_ids,
3477 vec![tid_bytes, tid_string, tid_string, tid_bytes],
3478 "branch selection for bytes/string union"
3479 );
3480 let s_child = u
3481 .child(tid_string)
3482 .as_any()
3483 .downcast_ref::<StringArray>()
3484 .unwrap();
3485 assert_eq!(s_child.len(), 2);
3486 assert_eq!(s_child.value(0), "hello");
3487 assert_eq!(s_child.value(1), "world");
3488 let b_child = u
3489 .child(tid_bytes)
3490 .as_any()
3491 .downcast_ref::<BinaryArray>()
3492 .unwrap();
3493 assert_eq!(b_child.len(), 2);
3494 assert_eq!(b_child.value(0), &[0x00, 0xFF, 0x7F]);
3495 assert_eq!(b_child.value(1), b""); let idx = schema.index_of("union_enum_records_array_map").unwrap();
3497 let u = batch
3498 .column(idx)
3499 .as_any()
3500 .downcast_ref::<UnionArray>()
3501 .expect("union_enum_records_array_map should be Union");
3502 let fields = match u.data_type() {
3503 DataType::Union(fields, _) => fields,
3504 other => panic!("expected Union, got {other:?}"),
3505 };
3506 let mut tid_enum: Option<i8> = None;
3507 let mut tid_rec_a: Option<i8> = None;
3508 let mut tid_rec_b: Option<i8> = None;
3509 let mut tid_array: Option<i8> = None;
3510 for (tid, f) in fields.iter() {
3511 match f.data_type() {
3512 DataType::Dictionary(_, _) => tid_enum = Some(tid),
3513 DataType::Struct(childs) => {
3514 if childs.len() == 2 && childs[0].name() == "a" && childs[1].name() == "b" {
3515 tid_rec_a = Some(tid);
3516 } else if childs.len() == 2
3517 && childs[0].name() == "x"
3518 && childs[1].name() == "y"
3519 {
3520 tid_rec_b = Some(tid);
3521 }
3522 }
3523 DataType::List(_) => tid_array = Some(tid),
3524 _ => {}
3525 }
3526 }
3527 let (tid_enum, tid_rec_a, tid_rec_b, tid_array) = (
3528 tid_enum.expect("enum child"),
3529 tid_rec_a.expect("RecA child"),
3530 tid_rec_b.expect("RecB child"),
3531 tid_array.expect("array<long> child"),
3532 );
3533 let type_ids: Vec<i8> = u.type_ids().iter().copied().collect();
3534 assert_eq!(
3535 type_ids,
3536 vec![tid_enum, tid_rec_a, tid_rec_b, tid_array],
3537 "branch selection for complex union"
3538 );
3539 let dict = u
3540 .child(tid_enum)
3541 .as_any()
3542 .downcast_ref::<DictionaryArray<Int32Type>>()
3543 .unwrap();
3544 assert_eq!(dict.len(), 1);
3545 assert!(dict.is_valid(0));
3546 let rec_a = u
3547 .child(tid_rec_a)
3548 .as_any()
3549 .downcast_ref::<StructArray>()
3550 .unwrap();
3551 assert_eq!(rec_a.len(), 1);
3552 let a_val = rec_a
3553 .column_by_name("a")
3554 .unwrap()
3555 .as_any()
3556 .downcast_ref::<Int32Array>()
3557 .unwrap();
3558 assert_eq!(a_val.value(0), 7);
3559 let b_val = rec_a
3560 .column_by_name("b")
3561 .unwrap()
3562 .as_any()
3563 .downcast_ref::<StringArray>()
3564 .unwrap();
3565 assert_eq!(b_val.value(0), "x");
3566 let rec_b = u
3568 .child(tid_rec_b)
3569 .as_any()
3570 .downcast_ref::<StructArray>()
3571 .unwrap();
3572 let x_val = rec_b
3573 .column_by_name("x")
3574 .unwrap()
3575 .as_any()
3576 .downcast_ref::<Int64Array>()
3577 .unwrap();
3578 assert_eq!(x_val.value(0), 123_456_789_i64);
3579 let y_val = rec_b
3580 .column_by_name("y")
3581 .unwrap()
3582 .as_any()
3583 .downcast_ref::<BinaryArray>()
3584 .unwrap();
3585 assert_eq!(y_val.value(0), &[0xFF, 0x00]);
3586 let arr = u
3587 .child(tid_array)
3588 .as_any()
3589 .downcast_ref::<ListArray>()
3590 .unwrap();
3591 assert_eq!(arr.len(), 1);
3592 let first_values = arr.value(0);
3593 let longs = first_values.as_any().downcast_ref::<Int64Array>().unwrap();
3594 assert_eq!(longs.len(), 3);
3595 assert_eq!(longs.value(0), 1);
3596 assert_eq!(longs.value(1), 2);
3597 assert_eq!(longs.value(2), 3);
3598 let idx = schema.index_of("union_date_or_fixed4").unwrap();
3599 let u = batch
3600 .column(idx)
3601 .as_any()
3602 .downcast_ref::<UnionArray>()
3603 .expect("union_date_or_fixed4 should be Union");
3604 let fields = match u.data_type() {
3605 DataType::Union(fields, _) => fields,
3606 other => panic!("expected Union, got {other:?}"),
3607 };
3608 let mut tid_date: Option<i8> = None;
3609 let mut tid_fixed: Option<i8> = None;
3610 for (tid, f) in fields.iter() {
3611 match f.data_type() {
3612 DataType::Date32 => tid_date = Some(tid),
3613 DataType::FixedSizeBinary(4) => tid_fixed = Some(tid),
3614 _ => {}
3615 }
3616 }
3617 let (tid_date, tid_fixed) = (tid_date.expect("date"), tid_fixed.expect("fixed(4)"));
3618 let type_ids: Vec<i8> = u.type_ids().iter().copied().collect();
3619 assert_eq!(
3620 type_ids,
3621 vec![tid_date, tid_fixed, tid_date, tid_fixed],
3622 "branch selection for date/fixed4 union"
3623 );
3624 let dates = u
3625 .child(tid_date)
3626 .as_any()
3627 .downcast_ref::<Date32Array>()
3628 .unwrap();
3629 assert_eq!(dates.len(), 2);
3630 assert_eq!(dates.value(0), 19_000); assert_eq!(dates.value(1), 0); let fixed = u
3633 .child(tid_fixed)
3634 .as_any()
3635 .downcast_ref::<FixedSizeBinaryArray>()
3636 .unwrap();
3637 assert_eq!(fixed.len(), 2);
3638 assert_eq!(fixed.value(0), b"ABCD");
3639 assert_eq!(fixed.value(1), &[0x00, 0x11, 0x22, 0x33]);
3640 }
3641
3642 #[test]
3643 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 = [Some(fx8_a)].into_iter();
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 (precision, scale) = match expected_dt {
5381 DataType::Decimal32(p, s)
5382 | DataType::Decimal64(p, s)
5383 | DataType::Decimal128(p, s)
5384 | DataType::Decimal256(p, s) => (p, s),
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: i128 = 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: bool = 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 fn comprehensive_e2e_resolution_test() {
8280 use serde_json::Value;
8281 use std::collections::HashMap;
8282
8283 fn make_comprehensive_reader_schema(path: &str) -> AvroSchema {
8296 fn set_type_string(f: &mut Value, new_ty: &str) {
8297 if let Some(ty) = f.get_mut("type") {
8298 match ty {
8299 Value::String(_) | Value::Object(_) => {
8300 *ty = Value::String(new_ty.to_string());
8301 }
8302 Value::Array(arr) => {
8303 for b in arr.iter_mut() {
8304 match b {
8305 Value::String(s) if s != "null" => {
8306 *b = Value::String(new_ty.to_string());
8307 break;
8308 }
8309 Value::Object(_) => {
8310 *b = Value::String(new_ty.to_string());
8311 break;
8312 }
8313 _ => {}
8314 }
8315 }
8316 }
8317 _ => {}
8318 }
8319 }
8320 }
8321 fn reverse_union_array(f: &mut Value) {
8322 if let Some(arr) = f.get_mut("type").and_then(|t| t.as_array_mut()) {
8323 arr.reverse();
8324 }
8325 }
8326 fn reverse_items_union(f: &mut Value) {
8327 if let Some(obj) = f.get_mut("type").and_then(|t| t.as_object_mut())
8328 && let Some(items) = obj.get_mut("items").and_then(|v| v.as_array_mut())
8329 {
8330 items.reverse();
8331 }
8332 }
8333 fn reverse_map_values_union(f: &mut Value) {
8334 if let Some(obj) = f.get_mut("type").and_then(|t| t.as_object_mut())
8335 && let Some(values) = obj.get_mut("values").and_then(|v| v.as_array_mut())
8336 {
8337 values.reverse();
8338 }
8339 }
8340 fn reverse_nested_union_in_record(f: &mut Value, field_name: &str) {
8341 if let Some(obj) = f.get_mut("type").and_then(|t| t.as_object_mut())
8342 && let Some(fields) = obj.get_mut("fields").and_then(|v| v.as_array_mut())
8343 {
8344 for ff in fields.iter_mut() {
8345 if ff.get("name").and_then(|n| n.as_str()) == Some(field_name)
8346 && let Some(ty) = ff.get_mut("type")
8347 && let Some(arr) = ty.as_array_mut()
8348 {
8349 arr.reverse();
8350 }
8351 }
8352 }
8353 }
8354 fn rename_nested_field_with_alias(f: &mut Value, old: &str, new: &str) {
8355 if let Some(obj) = f.get_mut("type").and_then(|t| t.as_object_mut())
8356 && let Some(fields) = obj.get_mut("fields").and_then(|v| v.as_array_mut())
8357 {
8358 for ff in fields.iter_mut() {
8359 if ff.get("name").and_then(|n| n.as_str()) == Some(old) {
8360 ff["name"] = Value::String(new.to_string());
8361 ff["aliases"] = Value::Array(vec![Value::String(old.to_string())]);
8362 }
8363 }
8364 }
8365 }
8366 let mut root = load_writer_schema_json(path);
8367 assert_eq!(root["type"], "record", "writer schema must be a record");
8368 let fields = root
8369 .get_mut("fields")
8370 .and_then(|f| f.as_array_mut())
8371 .expect("record has fields");
8372 for f in fields.iter_mut() {
8373 let Some(name) = f.get("name").and_then(|n| n.as_str()) else {
8374 continue;
8375 };
8376 match name {
8377 "id" => {
8379 f["name"] = Value::String("identifier".into());
8380 f["aliases"] = Value::Array(vec![Value::String("id".into())]);
8381 }
8382 "renamed_with_default" => {
8383 f["name"] = Value::String("old_count".into());
8384 f["aliases"] =
8385 Value::Array(vec![Value::String("renamed_with_default".into())]);
8386 }
8387 "count_i32" => set_type_string(f, "long"),
8389 "ratio_f32" => set_type_string(f, "double"),
8390 "opt_str_nullsecond" => reverse_union_array(f),
8392 "union_enum_record_array_map" => reverse_union_array(f),
8393 "union_date_or_fixed4" => reverse_union_array(f),
8394 "union_interval_or_string" => reverse_union_array(f),
8395 "union_uuid_or_fixed10" => reverse_union_array(f),
8396 "union_map_or_array_int" => reverse_union_array(f),
8397 "maybe_auth" => reverse_nested_union_in_record(f, "token"),
8398 "arr_union" => reverse_items_union(f),
8400 "map_union" => reverse_map_values_union(f),
8401 "address" => rename_nested_field_with_alias(f, "street", "street_name"),
8403 "person" => {
8405 if let Some(tobj) = f.get_mut("type").and_then(|t| t.as_object_mut()) {
8406 tobj.insert("name".to_string(), Value::String("Person".into()));
8407 tobj.insert(
8408 "namespace".to_string(),
8409 Value::String("com.example".into()),
8410 );
8411 tobj.insert(
8412 "aliases".into(),
8413 Value::Array(vec![
8414 Value::String("PersonV2".into()),
8415 Value::String("com.example.v2.PersonV2".into()),
8416 ]),
8417 );
8418 }
8419 }
8420 _ => {}
8421 }
8422 }
8423 fields.reverse();
8424 AvroSchema::new(root.to_string())
8425 }
8426
8427 let path = "test/data/comprehensive_e2e.avro";
8428 let reader_schema = make_comprehensive_reader_schema(path);
8429 let batch = read_alltypes_with_reader_schema(path, reader_schema.clone());
8430
8431 const UUID_EXT_KEY: &str = "ARROW:extension:name";
8432 const UUID_LOGICAL_KEY: &str = "logicalType";
8433
8434 let uuid_md_top: Option<arrow_schema::Metadata> = batch
8435 .schema()
8436 .field_with_name("uuid_str")
8437 .ok()
8438 .and_then(|f| {
8439 let md = f.metadata();
8440 let has_ext = md.get(UUID_EXT_KEY).is_some();
8441 let is_uuid_logical = md
8442 .get(UUID_LOGICAL_KEY)
8443 .map(|v| v.trim_matches('"') == "uuid")
8444 .unwrap_or(false);
8445 if has_ext || is_uuid_logical {
8446 Some(md.clone())
8447 } else {
8448 None
8449 }
8450 });
8451
8452 let uuid_md_union: Option<arrow_schema::Metadata> = batch
8453 .schema()
8454 .field_with_name("union_uuid_or_fixed10")
8455 .ok()
8456 .and_then(|f| match f.data_type() {
8457 DataType::Union(uf, _) => uf
8458 .iter()
8459 .find(|(_, child)| child.name() == "uuid")
8460 .and_then(|(_, child)| {
8461 let md = child.metadata();
8462 let has_ext = md.get(UUID_EXT_KEY).is_some();
8463 let is_uuid_logical = md
8464 .get(UUID_LOGICAL_KEY)
8465 .map(|v| v.trim_matches('"') == "uuid")
8466 .unwrap_or(false);
8467 if has_ext || is_uuid_logical {
8468 Some(md.clone())
8469 } else {
8470 None
8471 }
8472 }),
8473 _ => None,
8474 });
8475
8476 let add_uuid_ext_top = |f: Field| -> Field {
8477 if let Some(md) = &uuid_md_top {
8478 f.with_metadata(md.clone())
8479 } else {
8480 f
8481 }
8482 };
8483 let add_uuid_ext_union = |f: Field| -> Field {
8484 if let Some(md) = &uuid_md_union {
8485 f.with_metadata(md.clone())
8486 } else {
8487 f
8488 }
8489 };
8490
8491 #[inline]
8492 fn uuid16_from_str(s: &str) -> [u8; 16] {
8493 let mut out = [0u8; 16];
8494 let mut idx = 0usize;
8495 let mut hi: Option<u8> = None;
8496 for ch in s.chars() {
8497 if ch == '-' {
8498 continue;
8499 }
8500 let v = ch.to_digit(16).expect("invalid hex digit in UUID") as u8;
8501 if let Some(h) = hi {
8502 out[idx] = (h << 4) | v;
8503 idx += 1;
8504 hi = None;
8505 } else {
8506 hi = Some(v);
8507 }
8508 }
8509 assert_eq!(idx, 16, "UUID must decode to 16 bytes");
8510 out
8511 }
8512
8513 fn mk_dense_union(
8514 fields: &UnionFields,
8515 type_ids: Vec<i8>,
8516 offsets: Vec<i32>,
8517 provide: impl Fn(&Field) -> Option<ArrayRef>,
8518 ) -> ArrayRef {
8519 fn empty_child_for(dt: &DataType) -> Arc<dyn Array> {
8520 match dt {
8521 DataType::Null => Arc::new(NullArray::new(0)),
8522 DataType::Boolean => Arc::new(BooleanArray::from(Vec::<bool>::new())),
8523 DataType::Int32 => Arc::new(Int32Array::from(Vec::<i32>::new())),
8524 DataType::Int64 => Arc::new(Int64Array::from(Vec::<i64>::new())),
8525 DataType::Float32 => Arc::new(Float32Array::from(Vec::<f32>::new())),
8526 DataType::Float64 => Arc::new(Float64Array::from(Vec::<f64>::new())),
8527 DataType::Binary => Arc::new(BinaryArray::from(Vec::<&[u8]>::new())),
8528 DataType::Utf8 => Arc::new(StringArray::from(Vec::<&str>::new())),
8529 DataType::Date32 => Arc::new(Date32Array::from(Vec::<i32>::new())),
8530 DataType::Time32(arrow_schema::TimeUnit::Millisecond) => {
8531 Arc::new(Time32MillisecondArray::from(Vec::<i32>::new()))
8532 }
8533 DataType::Time64(arrow_schema::TimeUnit::Microsecond) => {
8534 Arc::new(Time64MicrosecondArray::from(Vec::<i64>::new()))
8535 }
8536 DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, tz) => {
8537 let a = TimestampMillisecondArray::from(Vec::<i64>::new());
8538 Arc::new(if let Some(tz) = tz {
8539 a.with_timezone(tz.clone())
8540 } else {
8541 a
8542 })
8543 }
8544 DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, tz) => {
8545 let a = TimestampMicrosecondArray::from(Vec::<i64>::new());
8546 Arc::new(if let Some(tz) = tz {
8547 a.with_timezone(tz.clone())
8548 } else {
8549 a
8550 })
8551 }
8552 DataType::Interval(IntervalUnit::MonthDayNano) => Arc::new(
8553 IntervalMonthDayNanoArray::from(Vec::<IntervalMonthDayNano>::new()),
8554 ),
8555 DataType::FixedSizeBinary(sz) => Arc::new(
8556 FixedSizeBinaryArray::try_from_sparse_iter_with_size(
8557 std::iter::empty::<Option<Vec<u8>>>(),
8558 *sz,
8559 )
8560 .unwrap(),
8561 ),
8562 DataType::Dictionary(_, _) => {
8563 let keys = Int32Array::from(Vec::<i32>::new());
8564 let values = Arc::new(StringArray::from(Vec::<&str>::new()));
8565 Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap())
8566 }
8567 DataType::Struct(fields) => {
8568 let children: Vec<ArrayRef> = fields
8569 .iter()
8570 .map(|f| empty_child_for(f.data_type()) as ArrayRef)
8571 .collect();
8572 Arc::new(StructArray::new(fields.clone(), children, None))
8573 }
8574 DataType::List(field) => {
8575 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0]));
8576 Arc::new(
8577 ListArray::try_new(
8578 field.clone(),
8579 offsets,
8580 empty_child_for(field.data_type()),
8581 None,
8582 )
8583 .unwrap(),
8584 )
8585 }
8586 DataType::Map(entry_field, is_sorted) => {
8587 let (key_field, val_field) = match entry_field.data_type() {
8588 DataType::Struct(fs) => (fs[0].clone(), fs[1].clone()),
8589 other => panic!("unexpected map entries type: {other:?}"),
8590 };
8591 let keys = StringArray::from(Vec::<&str>::new());
8592 let vals: ArrayRef = match val_field.data_type() {
8593 DataType::Null => Arc::new(NullArray::new(0)) as ArrayRef,
8594 DataType::Boolean => {
8595 Arc::new(BooleanArray::from(Vec::<bool>::new())) as ArrayRef
8596 }
8597 DataType::Int32 => {
8598 Arc::new(Int32Array::from(Vec::<i32>::new())) as ArrayRef
8599 }
8600 DataType::Int64 => {
8601 Arc::new(Int64Array::from(Vec::<i64>::new())) as ArrayRef
8602 }
8603 DataType::Float32 => {
8604 Arc::new(Float32Array::from(Vec::<f32>::new())) as ArrayRef
8605 }
8606 DataType::Float64 => {
8607 Arc::new(Float64Array::from(Vec::<f64>::new())) as ArrayRef
8608 }
8609 DataType::Utf8 => {
8610 Arc::new(StringArray::from(Vec::<&str>::new())) as ArrayRef
8611 }
8612 DataType::Binary => {
8613 Arc::new(BinaryArray::from(Vec::<&[u8]>::new())) as ArrayRef
8614 }
8615 DataType::Union(uf, _) => {
8616 let children: Vec<ArrayRef> = uf
8617 .iter()
8618 .map(|(_, f)| empty_child_for(f.data_type()))
8619 .collect();
8620 Arc::new(
8621 UnionArray::try_new(
8622 uf.clone(),
8623 ScalarBuffer::<i8>::from(Vec::<i8>::new()),
8624 Some(ScalarBuffer::<i32>::from(Vec::<i32>::new())),
8625 children,
8626 )
8627 .unwrap(),
8628 ) as ArrayRef
8629 }
8630 other => panic!("unsupported map value type: {other:?}"),
8631 };
8632 let entries = StructArray::new(
8633 Fields::from(vec![
8634 key_field.as_ref().clone(),
8635 val_field.as_ref().clone(),
8636 ]),
8637 vec![Arc::new(keys) as ArrayRef, vals],
8638 None,
8639 );
8640 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0]));
8641 Arc::new(MapArray::new(
8642 entry_field.clone(),
8643 offsets,
8644 entries,
8645 None,
8646 *is_sorted,
8647 ))
8648 }
8649 other => panic!("empty_child_for: unhandled type {other:?}"),
8650 }
8651 }
8652 let children: Vec<ArrayRef> = fields
8653 .iter()
8654 .map(|(_, f)| provide(f).unwrap_or_else(|| empty_child_for(f.data_type())))
8655 .collect();
8656 Arc::new(
8657 UnionArray::try_new(
8658 fields.clone(),
8659 ScalarBuffer::<i8>::from(type_ids),
8660 Some(ScalarBuffer::<i32>::from(offsets)),
8661 children,
8662 )
8663 .unwrap(),
8664 ) as ArrayRef
8665 }
8666 let date_a: i32 = 19_000; let time_ms_a: i32 = 12 * 3_600_000 + 34 * 60_000 + 56_000 + 789;
8668 let time_us_eod: i64 = 86_400_000_000 - 1;
8669 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;
8671 let dur_small = IntervalMonthDayNanoType::make_value(1, 2, 3_000_000_000);
8672 let dur_zero = IntervalMonthDayNanoType::make_value(0, 0, 0);
8673 let dur_large =
8674 IntervalMonthDayNanoType::make_value(12, 31, ((86_400_000 - 1) as i64) * 1_000_000);
8675 let dur_2years = IntervalMonthDayNanoType::make_value(24, 0, 0);
8676 let uuid1 = uuid16_from_str("fe7bc30b-4ce8-4c5e-b67c-2234a2d38e66");
8677 let uuid2 = uuid16_from_str("0826cc06-d2e3-4599-b4ad-af5fa6905cdb");
8678 let item_name = Field::LIST_FIELD_DEFAULT_NAME;
8679 let uf_tri = UnionFields::try_new(
8680 vec![0, 1, 2],
8681 vec![
8682 Field::new("int", DataType::Int32, false),
8683 Field::new("string", DataType::Utf8, false),
8684 Field::new("boolean", DataType::Boolean, false),
8685 ],
8686 )
8687 .unwrap();
8688 let uf_arr_items = UnionFields::try_new(
8689 vec![0, 1, 2],
8690 vec![
8691 Field::new("null", DataType::Null, false),
8692 Field::new("string", DataType::Utf8, false),
8693 Field::new("long", DataType::Int64, false),
8694 ],
8695 )
8696 .unwrap();
8697 let arr_items_field = Arc::new(Field::new(
8698 item_name,
8699 DataType::Union(uf_arr_items.clone(), UnionMode::Dense),
8700 true,
8701 ));
8702 let uf_map_vals = UnionFields::try_new(
8703 vec![0, 1, 2],
8704 vec![
8705 Field::new("string", DataType::Utf8, false),
8706 Field::new("double", DataType::Float64, false),
8707 Field::new("null", DataType::Null, false),
8708 ],
8709 )
8710 .unwrap();
8711 let map_entries_field = Arc::new(Field::new(
8712 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
8713 DataType::Struct(Fields::from(vec![
8714 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
8715 Field::new(
8716 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
8717 DataType::Union(uf_map_vals.clone(), UnionMode::Dense),
8718 true,
8719 ),
8720 ])),
8721 false,
8722 ));
8723 let mut enum_md_color = {
8725 let mut m = HashMap::<String, String>::new();
8726 m.insert(
8727 crate::schema::AVRO_ENUM_SYMBOLS_METADATA_KEY.to_string(),
8728 serde_json::to_string(&vec!["RED", "GREEN", "BLUE"]).unwrap(),
8729 );
8730 m
8731 };
8732 enum_md_color.insert(AVRO_NAME_METADATA_KEY.to_string(), "Color".to_string());
8733 enum_md_color.insert(
8734 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8735 "org.apache.arrow.avrotests.v1.types".to_string(),
8736 );
8737 let union_rec_a_fields = Fields::from(vec![
8738 Field::new("a", DataType::Int32, false),
8739 Field::new("b", DataType::Utf8, false),
8740 ]);
8741 let union_rec_b_fields = Fields::from(vec![
8742 Field::new("x", DataType::Int64, false),
8743 Field::new("y", DataType::Binary, false),
8744 ]);
8745 let union_map_entries = Arc::new(Field::new(
8746 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
8747 DataType::Struct(Fields::from(vec![
8748 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
8749 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, false),
8750 ])),
8751 false,
8752 ));
8753 let person_md = {
8754 let mut m = HashMap::<String, String>::new();
8755 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "Person".to_string());
8756 m.insert(
8757 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8758 "com.example".to_string(),
8759 );
8760 m
8761 };
8762 let maybe_auth_md = {
8763 let mut m = HashMap::<String, String>::new();
8764 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "MaybeAuth".to_string());
8765 m.insert(
8766 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8767 "org.apache.arrow.avrotests.v1.types".to_string(),
8768 );
8769 m
8770 };
8771 let address_md = {
8772 let mut m = HashMap::<String, String>::new();
8773 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "Address".to_string());
8774 m.insert(
8775 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8776 "org.apache.arrow.avrotests.v1.types".to_string(),
8777 );
8778 m
8779 };
8780 let rec_a_md = {
8781 let mut m = HashMap::<String, String>::new();
8782 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "RecA".to_string());
8783 m.insert(
8784 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8785 "org.apache.arrow.avrotests.v1.types".to_string(),
8786 );
8787 m
8788 };
8789 let rec_b_md = {
8790 let mut m = HashMap::<String, String>::new();
8791 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "RecB".to_string());
8792 m.insert(
8793 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8794 "org.apache.arrow.avrotests.v1.types".to_string(),
8795 );
8796 m
8797 };
8798 let uf_union_big = UnionFields::try_new(
8799 vec![0, 1, 2, 3, 4],
8800 vec![
8801 Field::new(
8802 "map",
8803 DataType::Map(union_map_entries.clone(), false),
8804 false,
8805 ),
8806 Field::new(
8807 "array",
8808 DataType::List(Arc::new(Field::new(item_name, DataType::Int64, false))),
8809 false,
8810 ),
8811 Field::new(
8812 "org.apache.arrow.avrotests.v1.types.RecB",
8813 DataType::Struct(union_rec_b_fields.clone()),
8814 false,
8815 )
8816 .with_metadata(rec_b_md.clone()),
8817 Field::new(
8818 "org.apache.arrow.avrotests.v1.types.RecA",
8819 DataType::Struct(union_rec_a_fields.clone()),
8820 false,
8821 )
8822 .with_metadata(rec_a_md.clone()),
8823 Field::new(
8824 "org.apache.arrow.avrotests.v1.types.Color",
8825 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
8826 false,
8827 )
8828 .with_metadata(enum_md_color.clone()),
8829 ],
8830 )
8831 .unwrap();
8832 let fx4_md = {
8833 let mut m = HashMap::<String, String>::new();
8834 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "Fx4".to_string());
8835 m.insert(
8836 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8837 "org.apache.arrow.avrotests.v1".to_string(),
8838 );
8839 m
8840 };
8841 let uf_date_fixed4 = UnionFields::try_new(
8842 vec![0, 1],
8843 vec![
8844 Field::new(
8845 "org.apache.arrow.avrotests.v1.Fx4",
8846 DataType::FixedSizeBinary(4),
8847 false,
8848 )
8849 .with_metadata(fx4_md.clone()),
8850 Field::new("date", DataType::Date32, false),
8851 ],
8852 )
8853 .unwrap();
8854 let dur12u_md = {
8855 let mut m = HashMap::<String, String>::new();
8856 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "Dur12U".to_string());
8857 m.insert(
8858 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8859 "org.apache.arrow.avrotests.v1".to_string(),
8860 );
8861 m
8862 };
8863 let uf_dur_or_str = UnionFields::try_new(
8864 vec![0, 1],
8865 vec![
8866 Field::new("string", DataType::Utf8, false),
8867 Field::new(
8868 "org.apache.arrow.avrotests.v1.Dur12U",
8869 DataType::Interval(arrow_schema::IntervalUnit::MonthDayNano),
8870 false,
8871 )
8872 .with_metadata(dur12u_md.clone()),
8873 ],
8874 )
8875 .unwrap();
8876 let fx10_md = {
8877 let mut m = HashMap::<String, String>::new();
8878 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "Fx10".to_string());
8879 m.insert(
8880 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8881 "org.apache.arrow.avrotests.v1".to_string(),
8882 );
8883 m
8884 };
8885 let uf_uuid_or_fx10 = UnionFields::try_new(
8886 vec![0, 1],
8887 vec![
8888 Field::new(
8889 "org.apache.arrow.avrotests.v1.Fx10",
8890 DataType::FixedSizeBinary(10),
8891 false,
8892 )
8893 .with_metadata(fx10_md.clone()),
8894 add_uuid_ext_union(Field::new("uuid", DataType::FixedSizeBinary(16), false)),
8895 ],
8896 )
8897 .unwrap();
8898 let uf_kv_val = UnionFields::try_new(
8899 vec![0, 1, 2],
8900 vec![
8901 Field::new("null", DataType::Null, false),
8902 Field::new("int", DataType::Int32, false),
8903 Field::new("long", DataType::Int64, false),
8904 ],
8905 )
8906 .unwrap();
8907 let kv_fields = Fields::from(vec![
8908 Field::new("key", DataType::Utf8, false),
8909 Field::new(
8910 "val",
8911 DataType::Union(uf_kv_val.clone(), UnionMode::Dense),
8912 true,
8913 ),
8914 ]);
8915 let kv_md = {
8916 let mut m = HashMap::<String, String>::new();
8917 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "KV".to_string());
8918 m.insert(
8919 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8920 "org.apache.arrow.avrotests.v1.types".to_string(),
8921 );
8922 m
8923 };
8924 let kv_item_field = Arc::new(
8925 Field::new(item_name, DataType::Struct(kv_fields.clone()), false).with_metadata(kv_md),
8926 );
8927 let map_int_entries = Arc::new(Field::new(
8928 Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
8929 DataType::Struct(Fields::from(vec![
8930 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
8931 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Int32, false),
8932 ])),
8933 false,
8934 ));
8935 let uf_map_or_array = UnionFields::try_new(
8936 vec![0, 1],
8937 vec![
8938 Field::new(
8939 "array",
8940 DataType::List(Arc::new(Field::new(item_name, DataType::Int32, false))),
8941 false,
8942 ),
8943 Field::new("map", DataType::Map(map_int_entries.clone(), false), false),
8944 ],
8945 )
8946 .unwrap();
8947 let mut enum_md_status = {
8948 let mut m = HashMap::<String, String>::new();
8949 m.insert(
8950 crate::schema::AVRO_ENUM_SYMBOLS_METADATA_KEY.to_string(),
8951 serde_json::to_string(&vec!["UNKNOWN", "NEW", "PROCESSING", "DONE"]).unwrap(),
8952 );
8953 m
8954 };
8955 enum_md_status.insert(AVRO_NAME_METADATA_KEY.to_string(), "Status".to_string());
8956 enum_md_status.insert(
8957 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8958 "org.apache.arrow.avrotests.v1.types".to_string(),
8959 );
8960 let mut dec20_md = HashMap::<String, String>::new();
8961 dec20_md.insert("precision".to_string(), "20".to_string());
8962 dec20_md.insert("scale".to_string(), "4".to_string());
8963 dec20_md.insert(AVRO_NAME_METADATA_KEY.to_string(), "DecFix20".to_string());
8964 dec20_md.insert(
8965 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8966 "org.apache.arrow.avrotests.v1.types".to_string(),
8967 );
8968 let mut dec10_md = HashMap::<String, String>::new();
8969 dec10_md.insert("precision".to_string(), "10".to_string());
8970 dec10_md.insert("scale".to_string(), "2".to_string());
8971 let fx16_top_md = {
8972 let mut m = HashMap::<String, String>::new();
8973 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "Fx16".to_string());
8974 m.insert(
8975 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8976 "org.apache.arrow.avrotests.v1.types".to_string(),
8977 );
8978 m
8979 };
8980 let dur12_top_md = {
8981 let mut m = HashMap::<String, String>::new();
8982 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "Dur12".to_string());
8983 m.insert(
8984 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8985 "org.apache.arrow.avrotests.v1.types".to_string(),
8986 );
8987 m
8988 };
8989 #[cfg(feature = "small_decimals")]
8990 let dec20_dt = DataType::Decimal128(20, 4);
8991 #[cfg(not(feature = "small_decimals"))]
8992 let dec20_dt = DataType::Decimal128(20, 4);
8993 #[cfg(feature = "small_decimals")]
8994 let dec10_dt = DataType::Decimal64(10, 2);
8995 #[cfg(not(feature = "small_decimals"))]
8996 let dec10_dt = DataType::Decimal128(10, 2);
8997 let fields: Vec<FieldRef> = vec![
8998 Arc::new(
8999 Field::new(
9000 "person",
9001 DataType::Struct(Fields::from(vec![
9002 Field::new("name", DataType::Utf8, false),
9003 Field::new("age", DataType::Int32, false),
9004 ])),
9005 false,
9006 )
9007 .with_metadata(person_md),
9008 ),
9009 Arc::new(Field::new("old_count", DataType::Int32, false)),
9010 Arc::new(Field::new(
9011 "union_map_or_array_int",
9012 DataType::Union(uf_map_or_array.clone(), UnionMode::Dense),
9013 false,
9014 )),
9015 Arc::new(Field::new(
9016 "array_records_with_union",
9017 DataType::List(kv_item_field.clone()),
9018 false,
9019 )),
9020 Arc::new(Field::new(
9021 "union_uuid_or_fixed10",
9022 DataType::Union(uf_uuid_or_fx10.clone(), UnionMode::Dense),
9023 false,
9024 )),
9025 Arc::new(Field::new(
9026 "union_interval_or_string",
9027 DataType::Union(uf_dur_or_str.clone(), UnionMode::Dense),
9028 false,
9029 )),
9030 Arc::new(Field::new(
9031 "union_date_or_fixed4",
9032 DataType::Union(uf_date_fixed4.clone(), UnionMode::Dense),
9033 false,
9034 )),
9035 Arc::new(Field::new(
9036 "union_enum_record_array_map",
9037 DataType::Union(uf_union_big.clone(), UnionMode::Dense),
9038 false,
9039 )),
9040 Arc::new(
9041 Field::new(
9042 "maybe_auth",
9043 DataType::Struct(Fields::from(vec![
9044 Field::new("user", DataType::Utf8, false),
9045 Field::new("token", DataType::Binary, true), ])),
9047 false,
9048 )
9049 .with_metadata(maybe_auth_md),
9050 ),
9051 Arc::new(
9052 Field::new(
9053 "address",
9054 DataType::Struct(Fields::from(vec![
9055 Field::new("street_name", DataType::Utf8, false),
9056 Field::new("zip", DataType::Int32, false),
9057 Field::new("country", DataType::Utf8, false),
9058 ])),
9059 false,
9060 )
9061 .with_metadata(address_md),
9062 ),
9063 Arc::new(Field::new(
9064 "map_union",
9065 DataType::Map(map_entries_field.clone(), false),
9066 false,
9067 )),
9068 Arc::new(Field::new(
9069 "arr_union",
9070 DataType::List(arr_items_field.clone()),
9071 false,
9072 )),
9073 Arc::new(
9074 Field::new(
9075 "status",
9076 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
9077 false,
9078 )
9079 .with_metadata(enum_md_status.clone()),
9080 ),
9081 Arc::new(
9082 Field::new(
9083 "interval_mdn",
9084 DataType::Interval(IntervalUnit::MonthDayNano),
9085 false,
9086 )
9087 .with_metadata(dur12_top_md.clone()),
9088 ),
9089 Arc::new(Field::new(
9090 "ts_micros_local",
9091 DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, None),
9092 false,
9093 )),
9094 Arc::new(Field::new(
9095 "ts_millis_local",
9096 DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, None),
9097 false,
9098 )),
9099 Arc::new(Field::new(
9100 "ts_micros_utc",
9101 DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, Some("+00:00".into())),
9102 false,
9103 )),
9104 Arc::new(Field::new(
9105 "ts_millis_utc",
9106 DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, Some("+00:00".into())),
9107 false,
9108 )),
9109 Arc::new(Field::new(
9110 "t_micros",
9111 DataType::Time64(arrow_schema::TimeUnit::Microsecond),
9112 false,
9113 )),
9114 Arc::new(Field::new(
9115 "t_millis",
9116 DataType::Time32(arrow_schema::TimeUnit::Millisecond),
9117 false,
9118 )),
9119 Arc::new(Field::new("d_date", DataType::Date32, false)),
9120 Arc::new(add_uuid_ext_top(Field::new(
9121 "uuid_str",
9122 DataType::FixedSizeBinary(16),
9123 false,
9124 ))),
9125 Arc::new(Field::new("dec_fix_s20_4", dec20_dt, false).with_metadata(dec20_md.clone())),
9126 Arc::new(
9127 Field::new("dec_bytes_s10_2", dec10_dt, false).with_metadata(dec10_md.clone()),
9128 ),
9129 Arc::new(
9130 Field::new("fx16_plain", DataType::FixedSizeBinary(16), false)
9131 .with_metadata(fx16_top_md.clone()),
9132 ),
9133 Arc::new(Field::new("raw_bytes", DataType::Binary, false)),
9134 Arc::new(Field::new("str_utf8", DataType::Utf8, false)),
9135 Arc::new(Field::new(
9136 "tri_union_prim",
9137 DataType::Union(uf_tri.clone(), UnionMode::Dense),
9138 false,
9139 )),
9140 Arc::new(Field::new("opt_str_nullsecond", DataType::Utf8, true)),
9141 Arc::new(Field::new("opt_i32_nullfirst", DataType::Int32, true)),
9142 Arc::new(Field::new("count_i64", DataType::Int64, false)),
9143 Arc::new(Field::new("count_i32", DataType::Int64, false)),
9144 Arc::new(Field::new("ratio_f64", DataType::Float64, false)),
9145 Arc::new(Field::new("ratio_f32", DataType::Float64, false)),
9146 Arc::new(Field::new("flag", DataType::Boolean, false)),
9147 Arc::new(Field::new("identifier", DataType::Int64, false)),
9148 ];
9149 let expected_schema = Arc::new(arrow_schema::Schema::new(Fields::from(fields)));
9150 let mut cols: Vec<ArrayRef> = vec![
9151 Arc::new(StructArray::new(
9152 match expected_schema
9153 .field_with_name("person")
9154 .unwrap()
9155 .data_type()
9156 {
9157 DataType::Struct(fs) => fs.clone(),
9158 _ => unreachable!(),
9159 },
9160 vec![
9161 Arc::new(StringArray::from(vec!["Alice", "Bob", "Carol", "Dave"])) as ArrayRef,
9162 Arc::new(Int32Array::from(vec![30, 0, 25, 41])) as ArrayRef,
9163 ],
9164 None,
9165 )) as ArrayRef,
9166 Arc::new(Int32Array::from(vec![100, 42, 7, 42])) as ArrayRef,
9167 ];
9168 {
9169 let map_child: ArrayRef = {
9170 let keys = StringArray::from(vec!["x", "y", "only"]);
9171 let vals = Int32Array::from(vec![1, 2, 10]);
9172 let entries = StructArray::new(
9173 Fields::from(vec![
9174 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
9175 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Int32, false),
9176 ]),
9177 vec![Arc::new(keys) as ArrayRef, Arc::new(vals) as ArrayRef],
9178 None,
9179 );
9180 let moff = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2, 3]));
9181 Arc::new(MapArray::new(
9182 map_int_entries.clone(),
9183 moff,
9184 entries,
9185 None,
9186 false,
9187 )) as ArrayRef
9188 };
9189 let list_child: ArrayRef = {
9190 let values = Int32Array::from(vec![1, 2, 3, 0]);
9191 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 4]));
9192 Arc::new(
9193 ListArray::try_new(
9194 Arc::new(Field::new(item_name, DataType::Int32, false)),
9195 offsets,
9196 Arc::new(values),
9197 None,
9198 )
9199 .unwrap(),
9200 ) as ArrayRef
9201 };
9202 let tids = vec![1, 0, 1, 0];
9203 let offs = vec![0, 0, 1, 1];
9204 let arr = mk_dense_union(&uf_map_or_array, tids, offs, |f| match f.name().as_str() {
9205 "array" => Some(list_child.clone()),
9206 "map" => Some(map_child.clone()),
9207 _ => None,
9208 });
9209 cols.push(arr);
9210 }
9211 {
9212 let keys = Arc::new(StringArray::from(vec!["k1", "k2", "k", "k3", "x"])) as ArrayRef;
9213 let type_ids = vec![1, 0, 2, 0, 1];
9214 let offsets = vec![0, 0, 0, 1, 1];
9215 let vals = mk_dense_union(&uf_kv_val, type_ids, offsets, |f| match f.data_type() {
9216 DataType::Int32 => Some(Arc::new(Int32Array::from(vec![5, -5])) as ArrayRef),
9217 DataType::Int64 => Some(Arc::new(Int64Array::from(vec![99i64])) as ArrayRef),
9218 DataType::Null => Some(Arc::new(NullArray::new(2)) as ArrayRef),
9219 _ => None,
9220 });
9221 let values_struct =
9222 Arc::new(StructArray::new(kv_fields.clone(), vec![keys, vals], None));
9223 let list_offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2, 3, 4, 5]));
9224 let arr = Arc::new(
9225 ListArray::try_new(kv_item_field.clone(), list_offsets, values_struct, None)
9226 .unwrap(),
9227 ) as ArrayRef;
9228 cols.push(arr);
9229 }
9230 {
9231 let type_ids = vec![1, 0, 1, 0]; let offs = vec![0, 0, 1, 1];
9233 let arr = mk_dense_union(&uf_uuid_or_fx10, type_ids, offs, |f| match f.data_type() {
9234 DataType::FixedSizeBinary(16) => {
9235 let it = [Some(uuid1), Some(uuid2)].into_iter();
9236 Some(Arc::new(
9237 FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 16).unwrap(),
9238 ) as ArrayRef)
9239 }
9240 DataType::FixedSizeBinary(10) => {
9241 let fx10_a = [0xAAu8; 10];
9242 let fx10_b = [0x00u8, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99];
9243 let it = [Some(fx10_a), Some(fx10_b)].into_iter();
9244 Some(Arc::new(
9245 FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 10).unwrap(),
9246 ) as ArrayRef)
9247 }
9248 _ => None,
9249 });
9250 cols.push(arr);
9251 }
9252 {
9253 let type_ids = vec![1, 0, 1, 0]; let offs = vec![0, 0, 1, 1];
9255 let arr = mk_dense_union(&uf_dur_or_str, type_ids, offs, |f| match f.data_type() {
9256 DataType::Interval(arrow_schema::IntervalUnit::MonthDayNano) => Some(Arc::new(
9257 IntervalMonthDayNanoArray::from(vec![dur_small, dur_large]),
9258 )
9259 as ArrayRef),
9260 DataType::Utf8 => Some(Arc::new(StringArray::from(vec![
9261 "duration-as-text",
9262 "iso-8601-period-P1Y",
9263 ])) as ArrayRef),
9264 _ => None,
9265 });
9266 cols.push(arr);
9267 }
9268 {
9269 let type_ids = vec![1, 0, 1, 0]; let offs = vec![0, 0, 1, 1];
9271 let arr = mk_dense_union(&uf_date_fixed4, type_ids, offs, |f| match f.data_type() {
9272 DataType::Date32 => Some(Arc::new(Date32Array::from(vec![date_a, 0])) as ArrayRef),
9273 DataType::FixedSizeBinary(4) => {
9274 let it = [Some(*b"\x00\x11\x22\x33"), Some(*b"ABCD")].into_iter();
9275 Some(Arc::new(
9276 FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 4).unwrap(),
9277 ) as ArrayRef)
9278 }
9279 _ => None,
9280 });
9281 cols.push(arr);
9282 }
9283 {
9284 let tids = vec![4, 3, 1, 0]; let offs = vec![0, 0, 0, 0];
9286 let arr = mk_dense_union(&uf_union_big, tids, offs, |f| match f.data_type() {
9287 DataType::Dictionary(_, _) => {
9288 let keys = Int32Array::from(vec![0i32]);
9289 let values =
9290 Arc::new(StringArray::from(vec!["RED", "GREEN", "BLUE"])) as ArrayRef;
9291 Some(
9292 Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap())
9293 as ArrayRef,
9294 )
9295 }
9296 DataType::Struct(fs) if fs == &union_rec_a_fields => {
9297 let a = Int32Array::from(vec![7]);
9298 let b = StringArray::from(vec!["rec"]);
9299 Some(Arc::new(StructArray::new(
9300 fs.clone(),
9301 vec![Arc::new(a) as ArrayRef, Arc::new(b) as ArrayRef],
9302 None,
9303 )) as ArrayRef)
9304 }
9305 DataType::List(_) => {
9306 let values = Int64Array::from(vec![1i64, 2, 3]);
9307 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3]));
9308 Some(Arc::new(
9309 ListArray::try_new(
9310 Arc::new(Field::new(item_name, DataType::Int64, false)),
9311 offsets,
9312 Arc::new(values),
9313 None,
9314 )
9315 .unwrap(),
9316 ) as ArrayRef)
9317 }
9318 DataType::Map(_, _) => {
9319 let keys = StringArray::from(vec!["k"]);
9320 let vals = StringArray::from(vec!["v"]);
9321 let entries = StructArray::new(
9322 Fields::from(vec![
9323 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
9324 Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, false),
9325 ]),
9326 vec![Arc::new(keys) as ArrayRef, Arc::new(vals) as ArrayRef],
9327 None,
9328 );
9329 let moff = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 1]));
9330 Some(Arc::new(MapArray::new(
9331 union_map_entries.clone(),
9332 moff,
9333 entries,
9334 None,
9335 false,
9336 )) as ArrayRef)
9337 }
9338 _ => None,
9339 });
9340 cols.push(arr);
9341 }
9342 {
9343 let fs = match expected_schema
9344 .field_with_name("maybe_auth")
9345 .unwrap()
9346 .data_type()
9347 {
9348 DataType::Struct(fs) => fs.clone(),
9349 _ => unreachable!(),
9350 };
9351 let user =
9352 Arc::new(StringArray::from(vec!["alice", "bob", "carol", "dave"])) as ArrayRef;
9353 let token_values: Vec<Option<&[u8]>> = vec![
9354 None,
9355 Some(b"\x01\x02\x03".as_ref()),
9356 None,
9357 Some(b"".as_ref()),
9358 ];
9359 let token = Arc::new(BinaryArray::from(token_values)) as ArrayRef;
9360 cols.push(Arc::new(StructArray::new(fs, vec![user, token], None)) as ArrayRef);
9361 }
9362 {
9363 let fs = match expected_schema
9364 .field_with_name("address")
9365 .unwrap()
9366 .data_type()
9367 {
9368 DataType::Struct(fs) => fs.clone(),
9369 _ => unreachable!(),
9370 };
9371 let street = Arc::new(StringArray::from(vec![
9372 "100 Main",
9373 "",
9374 "42 Galaxy Way",
9375 "End Ave",
9376 ])) as ArrayRef;
9377 let zip = Arc::new(Int32Array::from(vec![12345, 0, 42424, 1])) as ArrayRef;
9378 let country = Arc::new(StringArray::from(vec!["US", "CA", "US", "GB"])) as ArrayRef;
9379 cols.push(Arc::new(StructArray::new(fs, vec![street, zip, country], None)) as ArrayRef);
9380 }
9381 {
9382 let keys = StringArray::from(vec!["a", "b", "c", "neg", "pi", "ok"]);
9383 let moff = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 4, 4, 6]));
9384 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];
9388 let offsets = vec![0, 0, 0, 1, 2, 1];
9389 let pi_5dp = (std::f64::consts::PI * 100_000.0).trunc() / 100_000.0;
9390 let vals = mk_dense_union(&uf_map_vals, type_ids, offsets, |f| match f.data_type() {
9391 DataType::Float64 => {
9392 Some(Arc::new(Float64Array::from(vec![1.5f64, -0.5, pi_5dp])) as ArrayRef)
9393 }
9394 DataType::Utf8 => {
9395 Some(Arc::new(StringArray::from(vec!["yes", "true"])) as ArrayRef)
9396 }
9397 DataType::Null => Some(Arc::new(NullArray::new(1)) as ArrayRef),
9398 _ => None,
9399 });
9400 let entries = StructArray::new(
9401 Fields::from(vec![
9402 Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
9403 Field::new(
9404 Field::MAP_VALUE_FIELD_DEFAULT_NAME,
9405 DataType::Union(uf_map_vals.clone(), UnionMode::Dense),
9406 true,
9407 ),
9408 ]),
9409 vec![Arc::new(keys) as ArrayRef, vals],
9410 None,
9411 );
9412 let map = Arc::new(MapArray::new(
9413 map_entries_field.clone(),
9414 moff,
9415 entries,
9416 None,
9417 false,
9418 )) as ArrayRef;
9419 cols.push(map);
9420 }
9421 {
9422 let type_ids = vec![
9423 2, 1, 0, 2, 0, 1, 2, 2, 1, 0,
9424 2, ];
9426 let offsets = vec![0, 0, 0, 1, 1, 1, 2, 3, 2, 2, 4];
9427 let values =
9428 mk_dense_union(&uf_arr_items, type_ids, offsets, |f| match f.data_type() {
9429 DataType::Int64 => {
9430 Some(Arc::new(Int64Array::from(vec![1i64, -3, 0, -1, 0])) as ArrayRef)
9431 }
9432 DataType::Utf8 => {
9433 Some(Arc::new(StringArray::from(vec!["x", "z", "end"])) as ArrayRef)
9434 }
9435 DataType::Null => Some(Arc::new(NullArray::new(3)) as ArrayRef),
9436 _ => None,
9437 });
9438 let list_offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 4, 7, 8, 11]));
9439 let arr = Arc::new(
9440 ListArray::try_new(arr_items_field.clone(), list_offsets, values, None).unwrap(),
9441 ) as ArrayRef;
9442 cols.push(arr);
9443 }
9444 {
9445 let keys = Int32Array::from(vec![1, 2, 3, 0]); let values = Arc::new(StringArray::from(vec![
9447 "UNKNOWN",
9448 "NEW",
9449 "PROCESSING",
9450 "DONE",
9451 ])) as ArrayRef;
9452 let dict = DictionaryArray::<Int32Type>::try_new(keys, values).unwrap();
9453 cols.push(Arc::new(dict) as ArrayRef);
9454 }
9455 cols.push(Arc::new(IntervalMonthDayNanoArray::from(vec![
9456 dur_small, dur_zero, dur_large, dur_2years,
9457 ])) as ArrayRef);
9458 cols.push(Arc::new(TimestampMicrosecondArray::from(vec![
9459 ts_us_2024_01_01 + 123_456,
9460 0,
9461 ts_us_2024_01_01 + 101_112,
9462 987_654_321,
9463 ])) as ArrayRef);
9464 cols.push(Arc::new(TimestampMillisecondArray::from(vec![
9465 ts_ms_2024_01_01 + 86_400_000,
9466 0,
9467 ts_ms_2024_01_01 + 789,
9468 123_456_789,
9469 ])) as ArrayRef);
9470 {
9471 let a = TimestampMicrosecondArray::from(vec![
9472 ts_us_2024_01_01,
9473 1,
9474 ts_us_2024_01_01 + 456,
9475 0,
9476 ])
9477 .with_timezone("+00:00");
9478 cols.push(Arc::new(a) as ArrayRef);
9479 }
9480 {
9481 let a = TimestampMillisecondArray::from(vec![
9482 ts_ms_2024_01_01,
9483 -1,
9484 ts_ms_2024_01_01 + 123,
9485 0,
9486 ])
9487 .with_timezone("+00:00");
9488 cols.push(Arc::new(a) as ArrayRef);
9489 }
9490 cols.push(Arc::new(Time64MicrosecondArray::from(vec![
9491 time_us_eod,
9492 0,
9493 1,
9494 1_000_000,
9495 ])) as ArrayRef);
9496 cols.push(Arc::new(Time32MillisecondArray::from(vec![
9497 time_ms_a,
9498 0,
9499 1,
9500 86_400_000 - 1,
9501 ])) as ArrayRef);
9502 cols.push(Arc::new(Date32Array::from(vec![date_a, 0, 1, 365])) as ArrayRef);
9503 {
9504 let it = [Some(uuid1), Some(uuid2), Some(uuid1), Some(uuid2)].into_iter();
9505 cols.push(Arc::new(
9506 FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 16).unwrap(),
9507 ) as ArrayRef);
9508 }
9509 {
9510 #[cfg(feature = "small_decimals")]
9511 let arr = Arc::new(
9512 Decimal128Array::from_iter_values([1_234_567_891_234i128, -420_000i128, 0, -1i128])
9513 .with_precision_and_scale(20, 4)
9514 .unwrap(),
9515 ) as ArrayRef;
9516 #[cfg(not(feature = "small_decimals"))]
9517 let arr = Arc::new(
9518 Decimal128Array::from_iter_values([1_234_567_891_234i128, -420_000i128, 0, -1i128])
9519 .with_precision_and_scale(20, 4)
9520 .unwrap(),
9521 ) as ArrayRef;
9522 cols.push(arr);
9523 }
9524 {
9525 #[cfg(feature = "small_decimals")]
9526 let arr = Arc::new(
9527 Decimal64Array::from_iter_values([123456i64, -1, 0, 9_999_999_999i64])
9528 .with_precision_and_scale(10, 2)
9529 .unwrap(),
9530 ) as ArrayRef;
9531 #[cfg(not(feature = "small_decimals"))]
9532 let arr = Arc::new(
9533 Decimal128Array::from_iter_values([123456i128, -1, 0, 9_999_999_999i128])
9534 .with_precision_and_scale(10, 2)
9535 .unwrap(),
9536 ) as ArrayRef;
9537 cols.push(arr);
9538 }
9539 {
9540 let it = [
9541 Some(*b"0123456789ABCDEF"),
9542 Some([0u8; 16]),
9543 Some(*b"ABCDEFGHIJKLMNOP"),
9544 Some([0xAA; 16]),
9545 ]
9546 .into_iter();
9547 cols.push(Arc::new(
9548 FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 16).unwrap(),
9549 ) as ArrayRef);
9550 }
9551 cols.push(Arc::new(BinaryArray::from(vec![
9552 b"\x00\x01".as_ref(),
9553 b"".as_ref(),
9554 b"\xFF\x00".as_ref(),
9555 b"\x10\x20\x30\x40".as_ref(),
9556 ])) as ArrayRef);
9557 cols.push(Arc::new(StringArray::from(vec!["hello", "", "world", "✓ unicode"])) as ArrayRef);
9558 {
9559 let tids = vec![0, 1, 2, 1];
9560 let offs = vec![0, 0, 0, 1];
9561 let arr = mk_dense_union(&uf_tri, tids, offs, |f| match f.data_type() {
9562 DataType::Int32 => Some(Arc::new(Int32Array::from(vec![0])) as ArrayRef),
9563 DataType::Utf8 => Some(Arc::new(StringArray::from(vec!["hi", ""])) as ArrayRef),
9564 DataType::Boolean => Some(Arc::new(BooleanArray::from(vec![true])) as ArrayRef),
9565 _ => None,
9566 });
9567 cols.push(arr);
9568 }
9569 cols.push(Arc::new(StringArray::from(vec![
9570 Some("alpha"),
9571 None,
9572 Some("s3"),
9573 Some(""),
9574 ])) as ArrayRef);
9575 cols.push(Arc::new(Int32Array::from(vec![None, Some(42), None, Some(0)])) as ArrayRef);
9576 cols.push(Arc::new(Int64Array::from(vec![
9577 7_000_000_000i64,
9578 -2,
9579 0,
9580 -9_876_543_210i64,
9581 ])) as ArrayRef);
9582 cols.push(Arc::new(Int64Array::from(vec![7i64, -1, 0, 123])) as ArrayRef);
9583 cols.push(Arc::new(Float64Array::from(vec![2.5f64, -1.0, 7.0, -2.25])) as ArrayRef);
9584 cols.push(Arc::new(Float64Array::from(vec![1.25f64, -0.0, 3.5, 9.75])) as ArrayRef);
9585 cols.push(Arc::new(BooleanArray::from(vec![true, false, true, false])) as ArrayRef);
9586 cols.push(Arc::new(Int64Array::from(vec![1, 2, 3, 4])) as ArrayRef);
9587 let expected = RecordBatch::try_new(expected_schema, cols).unwrap();
9588 assert_eq!(
9589 expected, batch,
9590 "entire RecordBatch mismatch (schema, all columns, all rows)"
9591 );
9592 }
9593
9594 fn make_type_ref_ocf() -> Vec<u8> {
9604 use apache_avro::{Schema as ApacheSchema, Writer as ApacheWriter, types::Value};
9605 let schema_json = r#"{
9606 "type": "record", "name": "Root",
9607 "fields": [
9608 {"name": "ts", "type": {"type": "record", "name": "Timestamp", "fields": [
9609 {"name": "seconds", "type": "long"},
9610 {"name": "nanos", "type": "int"}
9611 ]}},
9612 {"name": "extra", "type": {"type": "record", "name": "Event", "fields": [
9613 {"name": "time", "type": "Timestamp"}
9614 ]}}
9615 ]
9616 }"#;
9617 let schema = ApacheSchema::parse_str(schema_json).expect("valid schema");
9618 let mut out = Vec::new();
9619 {
9620 let mut writer = ApacheWriter::new(&schema, &mut out);
9621 let ts_val = |s: i64, n: i32| {
9622 Value::Record(vec![
9623 ("seconds".into(), Value::Long(s)),
9624 ("nanos".into(), Value::Int(n)),
9625 ])
9626 };
9627 for (ts_s, ts_n, ex_s, ex_n) in [(1000i64, 100i32, -1i64, -1i32), (2000, 200, -2, -2)] {
9629 let row = Value::Record(vec![
9630 ("ts".into(), ts_val(ts_s, ts_n)),
9631 (
9632 "extra".into(),
9633 Value::Record(vec![("time".into(), ts_val(ex_s, ex_n))]),
9634 ),
9635 ]);
9636 writer.append_value_ref(&row).expect("append row");
9637 }
9638 writer.flush().expect("flush");
9639 }
9640 out
9641 }
9642
9643 #[test]
9652 fn test_nullable_reader_schema_vs_plain_writer_nested_struct() {
9653 let bytes = make_type_ref_ocf();
9654 let reader_schema = AvroSchema::new(
9655 r#"{"type":"record","name":"Root","fields":[
9656 {"name":"ts","type":["null",{"type":"record","name":"Timestamp","fields":[
9657 {"name":"seconds","type":["null","long"]},
9658 {"name":"nanos", "type":["null","int"]}
9659 ]}]}
9660 ]}"#
9661 .to_string(),
9662 );
9663 let mut reader = ReaderBuilder::new()
9664 .with_reader_schema(reader_schema)
9665 .build(Cursor::new(bytes))
9666 .expect("reader should build");
9667 let batch = reader
9668 .next()
9669 .expect("should have a batch")
9670 .expect("reading should succeed");
9671 assert_eq!(batch.num_rows(), 2);
9672 let ts = batch
9673 .column(0)
9674 .as_any()
9675 .downcast_ref::<StructArray>()
9676 .unwrap();
9677 let seconds = ts
9678 .column_by_name("seconds")
9679 .unwrap()
9680 .as_any()
9681 .downcast_ref::<Int64Array>()
9682 .unwrap();
9683 assert_eq!(seconds.value(0), 1000);
9684 assert_eq!(seconds.value(1), 2000);
9685 }
9686
9687 #[test]
9695 fn test_skipper_consumes_writer_only_struct_fields() {
9696 let bytes = make_type_ref_ocf();
9697 let reader_schema = AvroSchema::new(
9698 r#"{"type":"record","name":"Root","fields":[
9699 {"name":"ts","type":{"type":"record","name":"Timestamp","fields":[
9700 {"name":"seconds","type":"long"}
9701 ]}}
9702 ]}"#
9703 .to_string(),
9704 );
9705 let mut reader = ReaderBuilder::new()
9706 .with_reader_schema(reader_schema)
9707 .build(Cursor::new(bytes))
9708 .expect("reader should build");
9709 let batch = reader
9710 .next()
9711 .expect("should have a batch")
9712 .expect("Skipper must consume both seconds and nanos for extra.time");
9713 assert_eq!(batch.num_rows(), 2);
9714 let ts = batch
9715 .column(0)
9716 .as_any()
9717 .downcast_ref::<StructArray>()
9718 .unwrap();
9719 let seconds = ts
9720 .column_by_name("seconds")
9721 .unwrap()
9722 .as_any()
9723 .downcast_ref::<Int64Array>()
9724 .unwrap();
9725 assert_eq!(seconds.value(0), 1000);
9726 assert_eq!(seconds.value(1), 2000);
9727 }
9728
9729 #[test]
9738 fn test_skip_array_of_structs_uses_writer_schema_not_resolved() {
9739 use apache_avro::{Schema as ApacheSchema, Writer as ApacheWriter, types::Value};
9740 let schema_json = r#"{
9741 "type": "record", "name": "Root",
9742 "fields": [
9743 {"name": "ts", "type": {"type": "record", "name": "Timestamp", "fields": [
9744 {"name": "seconds", "type": "long"},
9745 {"name": "nanos", "type": "int"}
9746 ]}},
9747 {"name": "events", "type": {"type": "array", "items": {
9748 "type": "record", "name": "Event", "fields": [
9749 {"name": "time", "type": "Timestamp"}
9750 ]
9751 }}}
9752 ]
9753 }"#;
9754 let schema = ApacheSchema::parse_str(schema_json).expect("valid schema");
9755 let mut bytes = Vec::new();
9756 {
9757 let mut writer = ApacheWriter::new(&schema, &mut bytes);
9758 let ts_val = |s: i64, n: i32| {
9760 Value::Record(vec![
9761 ("seconds".into(), Value::Long(s)),
9762 ("nanos".into(), Value::Int(n)),
9763 ])
9764 };
9765 let row = Value::Record(vec![
9766 ("ts".into(), ts_val(100, 5)),
9767 (
9768 "events".into(),
9769 Value::Array(vec![Value::Record(vec![("time".into(), ts_val(200, 1))])]),
9770 ),
9771 ]);
9772 writer.append_value_ref(&row).expect("append row");
9773 writer.flush().expect("flush");
9774 }
9775
9776 let reader_schema = AvroSchema::new(
9778 r#"{"type":"record","name":"Root","fields":[
9779 {"name":"ts","type":["null",{"type":"record","name":"Timestamp","fields":[
9780 {"name":"seconds","type":["null","long"]},
9781 {"name":"nanos", "type":["null","int"]}
9782 ]}]}
9783 ]}"#
9784 .to_string(),
9785 );
9786 let mut reader = ReaderBuilder::new()
9787 .with_reader_schema(reader_schema)
9788 .build(Cursor::new(bytes))
9789 .expect("reader should build");
9790 let batch = reader
9791 .next()
9792 .expect("should have a batch")
9793 .expect("Skipper must consume all events bytes using writer field types");
9794 assert_eq!(batch.num_rows(), 1);
9795 let ts = batch
9796 .column(0)
9797 .as_any()
9798 .downcast_ref::<StructArray>()
9799 .unwrap();
9800 let seconds = ts
9801 .column_by_name("seconds")
9802 .unwrap()
9803 .as_any()
9804 .downcast_ref::<Int64Array>()
9805 .unwrap();
9806 assert_eq!(seconds.value(0), 100);
9807 }
9808}