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("value", DataType::Int64, false));
3067 let mut map_builder = MapBuilder::new(
3068 Some(builder::MapFieldNames {
3069 entry: "entries".to_string(),
3070 key: "key".to_string(),
3071 value: "value".to_string(),
3072 }),
3073 StringBuilder::new(),
3074 Int64Builder::new(),
3075 )
3076 .with_values_field(values_field);
3077 for _ in 0..num_rows {
3078 let (keys, vals) = map_builder.entries();
3079 keys.append_value("a");
3080 vals.append_value(1);
3081 keys.append_value("b");
3082 vals.append_value(2);
3083 map_builder.append(true).unwrap();
3084 }
3085 arrays.push(Arc::new(map_builder.finish()));
3086 let rec_fields: Fields = Fields::from(vec![
3087 Field::new("x", DataType::Int32, false),
3088 Field::new("y", DataType::Utf8, true),
3089 ]);
3090 let mut sb = StructBuilder::new(
3091 rec_fields.clone(),
3092 vec![
3093 Box::new(Int32Builder::new()),
3094 Box::new(StringBuilder::new()),
3095 ],
3096 );
3097 for _ in 0..num_rows {
3098 sb.field_builder::<Int32Builder>(0).unwrap().append_value(7);
3099 sb.field_builder::<StringBuilder>(1).unwrap().append_null();
3100 sb.append(true);
3101 }
3102 arrays.push(Arc::new(sb.finish()));
3103 arrays.push(Arc::new(Int32Array::from_iter(std::iter::repeat_n(
3104 None::<i32>,
3105 num_rows,
3106 ))));
3107 arrays.push(Arc::new(Int32Array::from_iter_values(std::iter::repeat_n(
3108 123, num_rows,
3109 ))));
3110 let expected = RecordBatch::try_new(actual.schema(), arrays).unwrap();
3111 assert_eq!(
3112 actual, expected,
3113 "defaults should materialize correctly for all fields"
3114 );
3115 }
3116
3117 #[test]
3118 fn test_schema_resolution_default_enum_invalid_symbol_errors() {
3119 let path = "test/data/skippable_types.avro";
3120 let bad_schema = make_reader_schema_with_default_fields(
3121 path,
3122 vec![serde_json::json!({
3123 "name":"bad_enum",
3124 "type":{"type":"enum","name":"E","symbols":["A","B","C"]},
3125 "default":"Z"
3126 })],
3127 );
3128 let file = File::open(path).unwrap();
3129 let res = ReaderBuilder::new()
3130 .with_reader_schema(bad_schema)
3131 .build(BufReader::new(file));
3132 let err = res.expect_err("expected enum default validation to fail");
3133 let msg = err.to_string();
3134 let lower_msg = msg.to_lowercase();
3135 assert!(
3136 lower_msg.contains("enum")
3137 && (lower_msg.contains("symbol") || lower_msg.contains("default")),
3138 "unexpected error: {msg}"
3139 );
3140 }
3141
3142 #[test]
3143 fn test_schema_resolution_default_fixed_size_mismatch_errors() {
3144 let path = "test/data/skippable_types.avro";
3145 let bad_schema = make_reader_schema_with_default_fields(
3146 path,
3147 vec![serde_json::json!({
3148 "name":"bad_fixed",
3149 "type":{"type":"fixed","name":"F","size":4},
3150 "default":"ABC"
3151 })],
3152 );
3153 let file = File::open(path).unwrap();
3154 let res = ReaderBuilder::new()
3155 .with_reader_schema(bad_schema)
3156 .build(BufReader::new(file));
3157 let err = res.expect_err("expected fixed default validation to fail");
3158 let msg = err.to_string();
3159 let lower_msg = msg.to_lowercase();
3160 assert!(
3161 lower_msg.contains("fixed")
3162 && (lower_msg.contains("size")
3163 || lower_msg.contains("length")
3164 || lower_msg.contains("does not match")),
3165 "unexpected error: {msg}"
3166 );
3167 }
3168
3169 #[test]
3170 fn test_timestamp_with_utc_tz() {
3171 let path = arrow_test_data("avro/alltypes_plain.avro");
3172 let reader_schema =
3173 make_reader_schema_with_selected_fields_in_order(&path, &["timestamp_col"]);
3174 let file = File::open(path).unwrap();
3175 let reader = ReaderBuilder::new()
3176 .with_batch_size(1024)
3177 .with_utf8_view(false)
3178 .with_reader_schema(reader_schema)
3179 .with_tz(Tz::Utc)
3180 .build(BufReader::new(file))
3181 .unwrap();
3182 let schema = reader.schema();
3183 let batches = reader.collect::<Result<Vec<_>, _>>().unwrap();
3184 let batch = arrow::compute::concat_batches(&schema, &batches).unwrap();
3185 let expected = RecordBatch::try_from_iter_with_nullable([(
3186 "timestamp_col",
3187 Arc::new(
3188 TimestampMicrosecondArray::from_iter_values([
3189 1235865600000000, 1235865660000000, 1238544000000000, 1238544060000000, 1233446400000000, 1233446460000000, 1230768000000000, 1230768060000000, ])
3198 .with_timezone("UTC"),
3199 ) as _,
3200 true,
3201 )])
3202 .unwrap();
3203 assert_eq!(batch, expected);
3204 }
3205
3206 #[test]
3207 #[cfg(feature = "snappy")]
3209 fn test_alltypes_skip_writer_fields_keep_double_only() {
3210 let file = arrow_test_data("avro/alltypes_plain.avro");
3211 let reader_schema =
3212 make_reader_schema_with_selected_fields_in_order(&file, &["double_col"]);
3213 let batch = read_alltypes_with_reader_schema(&file, reader_schema);
3214 let expected = RecordBatch::try_from_iter_with_nullable([(
3215 "double_col",
3216 Arc::new(Float64Array::from_iter_values(
3217 (0..8).map(|x| (x % 2) as f64 * 10.1),
3218 )) as _,
3219 true,
3220 )])
3221 .unwrap();
3222 assert_eq!(batch, expected);
3223 }
3224
3225 #[test]
3226 #[cfg(feature = "snappy")]
3228 fn test_alltypes_skip_writer_fields_reorder_and_skip_many() {
3229 let file = arrow_test_data("avro/alltypes_plain.avro");
3230 let reader_schema =
3231 make_reader_schema_with_selected_fields_in_order(&file, &["timestamp_col", "id"]);
3232 let batch = read_alltypes_with_reader_schema(&file, reader_schema);
3233 let expected = RecordBatch::try_from_iter_with_nullable([
3234 (
3235 "timestamp_col",
3236 Arc::new(
3237 TimestampMicrosecondArray::from_iter_values([
3238 1235865600000000, 1235865660000000, 1238544000000000, 1238544060000000, 1233446400000000, 1233446460000000, 1230768000000000, 1230768060000000, ])
3247 .with_timezone("+00:00"),
3248 ) as _,
3249 true,
3250 ),
3251 (
3252 "id",
3253 Arc::new(Int32Array::from(vec![4, 5, 6, 7, 2, 3, 0, 1])) as _,
3254 true,
3255 ),
3256 ])
3257 .unwrap();
3258 assert_eq!(batch, expected);
3259 }
3260
3261 #[test]
3262 fn test_skippable_types_project_each_field_individually() {
3263 let path = "test/data/skippable_types.avro";
3264 let full = read_file(path, 1024, false);
3265 let schema_full = full.schema();
3266 let num_rows = full.num_rows();
3267 let writer_json = load_writer_schema_json(path);
3268 assert_eq!(
3269 writer_json["type"], "record",
3270 "writer schema must be a record"
3271 );
3272 let fields_json = writer_json
3273 .get("fields")
3274 .and_then(|f| f.as_array())
3275 .expect("record has fields");
3276 assert_eq!(
3277 schema_full.fields().len(),
3278 fields_json.len(),
3279 "full read column count vs writer fields"
3280 );
3281 fn rebuild_list_array_with_element(
3282 col: &ArrayRef,
3283 new_elem: Arc<Field>,
3284 is_large: bool,
3285 ) -> ArrayRef {
3286 if is_large {
3287 let list = col
3288 .as_any()
3289 .downcast_ref::<LargeListArray>()
3290 .expect("expected LargeListArray");
3291 let offsets = list.offsets().clone();
3292 let values = list.values().clone();
3293 let validity = list.nulls().cloned();
3294 Arc::new(LargeListArray::try_new(new_elem, offsets, values, validity).unwrap())
3295 } else {
3296 let list = col
3297 .as_any()
3298 .downcast_ref::<ListArray>()
3299 .expect("expected ListArray");
3300 let offsets = list.offsets().clone();
3301 let values = list.values().clone();
3302 let validity = list.nulls().cloned();
3303 Arc::new(ListArray::try_new(new_elem, offsets, values, validity).unwrap())
3304 }
3305 }
3306 for (idx, f) in fields_json.iter().enumerate() {
3307 let name = f
3308 .get("name")
3309 .and_then(|n| n.as_str())
3310 .unwrap_or_else(|| panic!("field at index {idx} has no name"));
3311 let reader_schema = make_reader_schema_with_selected_fields_in_order(path, &[name]);
3312 let projected = read_alltypes_with_reader_schema(path, reader_schema);
3313 assert_eq!(
3314 projected.num_columns(),
3315 1,
3316 "projected batch should contain exactly the selected column '{name}'"
3317 );
3318 assert_eq!(
3319 projected.num_rows(),
3320 num_rows,
3321 "row count mismatch for projected column '{name}'"
3322 );
3323 let col_full = full.column(idx).clone();
3324 let full_field = schema_full.field(idx).as_ref().clone();
3325 let proj_field_ref = projected.schema().field(0).clone();
3326 let proj_field = proj_field_ref.as_ref();
3327 let top_meta = proj_field.metadata().clone();
3328 let (expected_field_ref, expected_col): (Arc<Field>, ArrayRef) =
3329 match (full_field.data_type(), proj_field.data_type()) {
3330 (&DataType::List(_), DataType::List(proj_elem)) => {
3331 let new_col =
3332 rebuild_list_array_with_element(&col_full, proj_elem.clone(), false);
3333 let nf = Field::new(
3334 full_field.name().clone(),
3335 proj_field.data_type().clone(),
3336 full_field.is_nullable(),
3337 )
3338 .with_metadata(top_meta);
3339 (Arc::new(nf), new_col)
3340 }
3341 (&DataType::LargeList(_), DataType::LargeList(proj_elem)) => {
3342 let new_col =
3343 rebuild_list_array_with_element(&col_full, proj_elem.clone(), true);
3344 let nf = Field::new(
3345 full_field.name().clone(),
3346 proj_field.data_type().clone(),
3347 full_field.is_nullable(),
3348 )
3349 .with_metadata(top_meta);
3350 (Arc::new(nf), new_col)
3351 }
3352 _ => {
3353 let nf = full_field.with_metadata(top_meta);
3354 (Arc::new(nf), col_full)
3355 }
3356 };
3357
3358 let expected = RecordBatch::try_new(
3359 Arc::new(Schema::new(vec![expected_field_ref])),
3360 vec![expected_col],
3361 )
3362 .unwrap();
3363 assert_eq!(
3364 projected, expected,
3365 "projected column '{name}' mismatch vs full read column"
3366 );
3367 }
3368 }
3369
3370 #[test]
3371 fn test_union_fields_avro_nullable_and_general_unions() {
3372 let path = "test/data/union_fields.avro";
3373 let batch = read_file(path, 1024, false);
3374 let schema = batch.schema();
3375 let idx = schema.index_of("nullable_int_nullfirst").unwrap();
3376 let a = batch.column(idx).as_primitive::<Int32Type>();
3377 assert_eq!(a.len(), 4);
3378 assert!(a.is_null(0));
3379 assert_eq!(a.value(1), 42);
3380 assert!(a.is_null(2));
3381 assert_eq!(a.value(3), 0);
3382 let idx = schema.index_of("nullable_string_nullsecond").unwrap();
3383 let s = batch
3384 .column(idx)
3385 .as_any()
3386 .downcast_ref::<StringArray>()
3387 .expect("nullable_string_nullsecond should be Utf8");
3388 assert_eq!(s.len(), 4);
3389 assert_eq!(s.value(0), "s1");
3390 assert!(s.is_null(1));
3391 assert_eq!(s.value(2), "s3");
3392 assert!(s.is_valid(3)); assert_eq!(s.value(3), "");
3394 let idx = schema.index_of("union_prim").unwrap();
3395 let u = batch
3396 .column(idx)
3397 .as_any()
3398 .downcast_ref::<UnionArray>()
3399 .expect("union_prim should be Union");
3400 let fields = match u.data_type() {
3401 DataType::Union(fields, mode) => {
3402 assert!(matches!(mode, UnionMode::Dense), "expect dense unions");
3403 fields
3404 }
3405 other => panic!("expected Union, got {other:?}"),
3406 };
3407 let tid_by_name = |name: &str| -> i8 {
3408 for (tid, f) in fields.iter() {
3409 if f.name() == name {
3410 return tid;
3411 }
3412 }
3413 panic!("union child '{name}' not found");
3414 };
3415 let expected_type_ids = vec![
3416 tid_by_name("long"),
3417 tid_by_name("int"),
3418 tid_by_name("float"),
3419 tid_by_name("double"),
3420 ];
3421 let type_ids: Vec<i8> = u.type_ids().iter().copied().collect();
3422 assert_eq!(
3423 type_ids, expected_type_ids,
3424 "branch selection for union_prim rows"
3425 );
3426 let longs = u
3427 .child(tid_by_name("long"))
3428 .as_any()
3429 .downcast_ref::<Int64Array>()
3430 .unwrap();
3431 assert_eq!(longs.len(), 1);
3432 let ints = u
3433 .child(tid_by_name("int"))
3434 .as_any()
3435 .downcast_ref::<Int32Array>()
3436 .unwrap();
3437 assert_eq!(ints.len(), 1);
3438 let floats = u
3439 .child(tid_by_name("float"))
3440 .as_any()
3441 .downcast_ref::<Float32Array>()
3442 .unwrap();
3443 assert_eq!(floats.len(), 1);
3444 let doubles = u
3445 .child(tid_by_name("double"))
3446 .as_any()
3447 .downcast_ref::<Float64Array>()
3448 .unwrap();
3449 assert_eq!(doubles.len(), 1);
3450 let idx = schema.index_of("union_bytes_vs_string").unwrap();
3451 let u = batch
3452 .column(idx)
3453 .as_any()
3454 .downcast_ref::<UnionArray>()
3455 .expect("union_bytes_vs_string should be Union");
3456 let fields = match u.data_type() {
3457 DataType::Union(fields, _) => fields,
3458 other => panic!("expected Union, got {other:?}"),
3459 };
3460 let tid_by_name = |name: &str| -> i8 {
3461 for (tid, f) in fields.iter() {
3462 if f.name() == name {
3463 return tid;
3464 }
3465 }
3466 panic!("union child '{name}' not found");
3467 };
3468 let tid_bytes = tid_by_name("bytes");
3469 let tid_string = tid_by_name("string");
3470 let type_ids: Vec<i8> = u.type_ids().iter().copied().collect();
3471 assert_eq!(
3472 type_ids,
3473 vec![tid_bytes, tid_string, tid_string, tid_bytes],
3474 "branch selection for bytes/string union"
3475 );
3476 let s_child = u
3477 .child(tid_string)
3478 .as_any()
3479 .downcast_ref::<StringArray>()
3480 .unwrap();
3481 assert_eq!(s_child.len(), 2);
3482 assert_eq!(s_child.value(0), "hello");
3483 assert_eq!(s_child.value(1), "world");
3484 let b_child = u
3485 .child(tid_bytes)
3486 .as_any()
3487 .downcast_ref::<BinaryArray>()
3488 .unwrap();
3489 assert_eq!(b_child.len(), 2);
3490 assert_eq!(b_child.value(0), &[0x00, 0xFF, 0x7F]);
3491 assert_eq!(b_child.value(1), b""); let idx = schema.index_of("union_enum_records_array_map").unwrap();
3493 let u = batch
3494 .column(idx)
3495 .as_any()
3496 .downcast_ref::<UnionArray>()
3497 .expect("union_enum_records_array_map should be Union");
3498 let fields = match u.data_type() {
3499 DataType::Union(fields, _) => fields,
3500 other => panic!("expected Union, got {other:?}"),
3501 };
3502 let mut tid_enum: Option<i8> = None;
3503 let mut tid_rec_a: Option<i8> = None;
3504 let mut tid_rec_b: Option<i8> = None;
3505 let mut tid_array: Option<i8> = None;
3506 for (tid, f) in fields.iter() {
3507 match f.data_type() {
3508 DataType::Dictionary(_, _) => tid_enum = Some(tid),
3509 DataType::Struct(childs) => {
3510 if childs.len() == 2 && childs[0].name() == "a" && childs[1].name() == "b" {
3511 tid_rec_a = Some(tid);
3512 } else if childs.len() == 2
3513 && childs[0].name() == "x"
3514 && childs[1].name() == "y"
3515 {
3516 tid_rec_b = Some(tid);
3517 }
3518 }
3519 DataType::List(_) => tid_array = Some(tid),
3520 _ => {}
3521 }
3522 }
3523 let (tid_enum, tid_rec_a, tid_rec_b, tid_array) = (
3524 tid_enum.expect("enum child"),
3525 tid_rec_a.expect("RecA child"),
3526 tid_rec_b.expect("RecB child"),
3527 tid_array.expect("array<long> child"),
3528 );
3529 let type_ids: Vec<i8> = u.type_ids().iter().copied().collect();
3530 assert_eq!(
3531 type_ids,
3532 vec![tid_enum, tid_rec_a, tid_rec_b, tid_array],
3533 "branch selection for complex union"
3534 );
3535 let dict = u
3536 .child(tid_enum)
3537 .as_any()
3538 .downcast_ref::<DictionaryArray<Int32Type>>()
3539 .unwrap();
3540 assert_eq!(dict.len(), 1);
3541 assert!(dict.is_valid(0));
3542 let rec_a = u
3543 .child(tid_rec_a)
3544 .as_any()
3545 .downcast_ref::<StructArray>()
3546 .unwrap();
3547 assert_eq!(rec_a.len(), 1);
3548 let a_val = rec_a
3549 .column_by_name("a")
3550 .unwrap()
3551 .as_any()
3552 .downcast_ref::<Int32Array>()
3553 .unwrap();
3554 assert_eq!(a_val.value(0), 7);
3555 let b_val = rec_a
3556 .column_by_name("b")
3557 .unwrap()
3558 .as_any()
3559 .downcast_ref::<StringArray>()
3560 .unwrap();
3561 assert_eq!(b_val.value(0), "x");
3562 let rec_b = u
3564 .child(tid_rec_b)
3565 .as_any()
3566 .downcast_ref::<StructArray>()
3567 .unwrap();
3568 let x_val = rec_b
3569 .column_by_name("x")
3570 .unwrap()
3571 .as_any()
3572 .downcast_ref::<Int64Array>()
3573 .unwrap();
3574 assert_eq!(x_val.value(0), 123_456_789_i64);
3575 let y_val = rec_b
3576 .column_by_name("y")
3577 .unwrap()
3578 .as_any()
3579 .downcast_ref::<BinaryArray>()
3580 .unwrap();
3581 assert_eq!(y_val.value(0), &[0xFF, 0x00]);
3582 let arr = u
3583 .child(tid_array)
3584 .as_any()
3585 .downcast_ref::<ListArray>()
3586 .unwrap();
3587 assert_eq!(arr.len(), 1);
3588 let first_values = arr.value(0);
3589 let longs = first_values.as_any().downcast_ref::<Int64Array>().unwrap();
3590 assert_eq!(longs.len(), 3);
3591 assert_eq!(longs.value(0), 1);
3592 assert_eq!(longs.value(1), 2);
3593 assert_eq!(longs.value(2), 3);
3594 let idx = schema.index_of("union_date_or_fixed4").unwrap();
3595 let u = batch
3596 .column(idx)
3597 .as_any()
3598 .downcast_ref::<UnionArray>()
3599 .expect("union_date_or_fixed4 should be Union");
3600 let fields = match u.data_type() {
3601 DataType::Union(fields, _) => fields,
3602 other => panic!("expected Union, got {other:?}"),
3603 };
3604 let mut tid_date: Option<i8> = None;
3605 let mut tid_fixed: Option<i8> = None;
3606 for (tid, f) in fields.iter() {
3607 match f.data_type() {
3608 DataType::Date32 => tid_date = Some(tid),
3609 DataType::FixedSizeBinary(4) => tid_fixed = Some(tid),
3610 _ => {}
3611 }
3612 }
3613 let (tid_date, tid_fixed) = (tid_date.expect("date"), tid_fixed.expect("fixed(4)"));
3614 let type_ids: Vec<i8> = u.type_ids().iter().copied().collect();
3615 assert_eq!(
3616 type_ids,
3617 vec![tid_date, tid_fixed, tid_date, tid_fixed],
3618 "branch selection for date/fixed4 union"
3619 );
3620 let dates = u
3621 .child(tid_date)
3622 .as_any()
3623 .downcast_ref::<Date32Array>()
3624 .unwrap();
3625 assert_eq!(dates.len(), 2);
3626 assert_eq!(dates.value(0), 19_000); assert_eq!(dates.value(1), 0); let fixed = u
3629 .child(tid_fixed)
3630 .as_any()
3631 .downcast_ref::<FixedSizeBinaryArray>()
3632 .unwrap();
3633 assert_eq!(fixed.len(), 2);
3634 assert_eq!(fixed.value(0), b"ABCD");
3635 assert_eq!(fixed.value(1), &[0x00, 0x11, 0x22, 0x33]);
3636 }
3637
3638 #[test]
3639 fn test_union_schema_resolution_all_type_combinations() {
3640 let path = "test/data/union_fields.avro";
3641 let baseline = read_file(path, 1024, false);
3642 let baseline_schema = baseline.schema();
3643 let mut root = load_writer_schema_json(path);
3644 assert_eq!(root["type"], "record", "writer schema must be a record");
3645 let fields = root
3646 .get_mut("fields")
3647 .and_then(|f| f.as_array_mut())
3648 .expect("record has fields");
3649 fn is_named_type(obj: &Value, ty: &str, nm: &str) -> bool {
3650 obj.get("type").and_then(|v| v.as_str()) == Some(ty)
3651 && obj.get("name").and_then(|v| v.as_str()) == Some(nm)
3652 }
3653 fn is_logical(obj: &Value, prim: &str, lt: &str) -> bool {
3654 obj.get("type").and_then(|v| v.as_str()) == Some(prim)
3655 && obj.get("logicalType").and_then(|v| v.as_str()) == Some(lt)
3656 }
3657 fn find_first(arr: &[Value], pred: impl Fn(&Value) -> bool) -> Option<Value> {
3658 arr.iter().find(|v| pred(v)).cloned()
3659 }
3660 fn prim(s: &str) -> Value {
3661 Value::String(s.to_string())
3662 }
3663 for f in fields.iter_mut() {
3664 let Some(name) = f.get("name").and_then(|n| n.as_str()) else {
3665 continue;
3666 };
3667 match name {
3668 "nullable_int_nullfirst" => {
3670 f["type"] = json!(["int", "null"]);
3671 }
3672 "nullable_string_nullsecond" => {
3673 f["type"] = json!(["null", "string"]);
3674 }
3675 "union_prim" => {
3676 let orig = f["type"].as_array().unwrap().clone();
3677 let long = prim("long");
3678 let double = prim("double");
3679 let string = prim("string");
3680 let bytes = prim("bytes");
3681 let boolean = prim("boolean");
3682 assert!(orig.contains(&long));
3683 assert!(orig.contains(&double));
3684 assert!(orig.contains(&string));
3685 assert!(orig.contains(&bytes));
3686 assert!(orig.contains(&boolean));
3687 f["type"] = json!([long, double, string, bytes, boolean]);
3688 }
3689 "union_bytes_vs_string" => {
3690 f["type"] = json!(["string", "bytes"]);
3691 }
3692 "union_fixed_dur_decfix" => {
3693 let orig = f["type"].as_array().unwrap().clone();
3694 let fx8 = find_first(&orig, |o| is_named_type(o, "fixed", "Fx8")).unwrap();
3695 let dur12 = find_first(&orig, |o| is_named_type(o, "fixed", "Dur12")).unwrap();
3696 let decfix16 =
3697 find_first(&orig, |o| is_named_type(o, "fixed", "DecFix16")).unwrap();
3698 f["type"] = json!([decfix16, dur12, fx8]);
3699 }
3700 "union_enum_records_array_map" => {
3701 let orig = f["type"].as_array().unwrap().clone();
3702 let enum_color = find_first(&orig, |o| {
3703 o.get("type").and_then(|v| v.as_str()) == Some("enum")
3704 })
3705 .unwrap();
3706 let rec_a = find_first(&orig, |o| is_named_type(o, "record", "RecA")).unwrap();
3707 let rec_b = find_first(&orig, |o| is_named_type(o, "record", "RecB")).unwrap();
3708 let arr = find_first(&orig, |o| {
3709 o.get("type").and_then(|v| v.as_str()) == Some("array")
3710 })
3711 .unwrap();
3712 let map = find_first(&orig, |o| {
3713 o.get("type").and_then(|v| v.as_str()) == Some("map")
3714 })
3715 .unwrap();
3716 f["type"] = json!([arr, map, rec_b, rec_a, enum_color]);
3717 }
3718 "union_date_or_fixed4" => {
3719 let orig = f["type"].as_array().unwrap().clone();
3720 let date = find_first(&orig, |o| is_logical(o, "int", "date")).unwrap();
3721 let fx4 = find_first(&orig, |o| is_named_type(o, "fixed", "Fx4")).unwrap();
3722 f["type"] = json!([fx4, date]);
3723 }
3724 "union_time_millis_or_enum" => {
3725 let orig = f["type"].as_array().unwrap().clone();
3726 let time_ms =
3727 find_first(&orig, |o| is_logical(o, "int", "time-millis")).unwrap();
3728 let en = find_first(&orig, |o| {
3729 o.get("type").and_then(|v| v.as_str()) == Some("enum")
3730 })
3731 .unwrap();
3732 f["type"] = json!([en, time_ms]);
3733 }
3734 "union_time_micros_or_string" => {
3735 let orig = f["type"].as_array().unwrap().clone();
3736 let time_us =
3737 find_first(&orig, |o| is_logical(o, "long", "time-micros")).unwrap();
3738 f["type"] = json!(["string", time_us]);
3739 }
3740 "union_ts_millis_utc_or_array" => {
3741 let orig = f["type"].as_array().unwrap().clone();
3742 let ts_ms =
3743 find_first(&orig, |o| is_logical(o, "long", "timestamp-millis")).unwrap();
3744 let arr = find_first(&orig, |o| {
3745 o.get("type").and_then(|v| v.as_str()) == Some("array")
3746 })
3747 .unwrap();
3748 f["type"] = json!([arr, ts_ms]);
3749 }
3750 "union_ts_micros_local_or_bytes" => {
3751 let orig = f["type"].as_array().unwrap().clone();
3752 let lts_us =
3753 find_first(&orig, |o| is_logical(o, "long", "local-timestamp-micros"))
3754 .unwrap();
3755 f["type"] = json!(["bytes", lts_us]);
3756 }
3757 "union_uuid_or_fixed10" => {
3758 let orig = f["type"].as_array().unwrap().clone();
3759 let uuid = find_first(&orig, |o| is_logical(o, "string", "uuid")).unwrap();
3760 let fx10 = find_first(&orig, |o| is_named_type(o, "fixed", "Fx10")).unwrap();
3761 f["type"] = json!([fx10, uuid]);
3762 }
3763 "union_dec_bytes_or_dec_fixed" => {
3764 let orig = f["type"].as_array().unwrap().clone();
3765 let dec_bytes = find_first(&orig, |o| {
3766 o.get("type").and_then(|v| v.as_str()) == Some("bytes")
3767 && o.get("logicalType").and_then(|v| v.as_str()) == Some("decimal")
3768 })
3769 .unwrap();
3770 let dec_fix = find_first(&orig, |o| {
3771 is_named_type(o, "fixed", "DecFix20")
3772 && o.get("logicalType").and_then(|v| v.as_str()) == Some("decimal")
3773 })
3774 .unwrap();
3775 f["type"] = json!([dec_fix, dec_bytes]);
3776 }
3777 "union_null_bytes_string" => {
3778 f["type"] = json!(["bytes", "string", "null"]);
3779 }
3780 "array_of_union" => {
3781 let obj = f
3782 .get_mut("type")
3783 .expect("array type")
3784 .as_object_mut()
3785 .unwrap();
3786 obj.insert("items".to_string(), json!(["string", "long"]));
3787 }
3788 "map_of_union" => {
3789 let obj = f
3790 .get_mut("type")
3791 .expect("map type")
3792 .as_object_mut()
3793 .unwrap();
3794 obj.insert("values".to_string(), json!(["double", "null"]));
3795 }
3796 "record_with_union_field" => {
3797 let rec = f
3798 .get_mut("type")
3799 .expect("record type")
3800 .as_object_mut()
3801 .unwrap();
3802 let rec_fields = rec.get_mut("fields").unwrap().as_array_mut().unwrap();
3803 let mut found = false;
3804 for rf in rec_fields.iter_mut() {
3805 if rf.get("name").and_then(|v| v.as_str()) == Some("u") {
3806 rf["type"] = json!(["string", "long"]); found = true;
3808 break;
3809 }
3810 }
3811 assert!(found, "field 'u' expected in HasUnion");
3812 }
3813 "union_ts_micros_utc_or_map" => {
3814 let orig = f["type"].as_array().unwrap().clone();
3815 let ts_us =
3816 find_first(&orig, |o| is_logical(o, "long", "timestamp-micros")).unwrap();
3817 let map = find_first(&orig, |o| {
3818 o.get("type").and_then(|v| v.as_str()) == Some("map")
3819 })
3820 .unwrap();
3821 f["type"] = json!([map, ts_us]);
3822 }
3823 "union_ts_millis_local_or_string" => {
3824 let orig = f["type"].as_array().unwrap().clone();
3825 let lts_ms =
3826 find_first(&orig, |o| is_logical(o, "long", "local-timestamp-millis"))
3827 .unwrap();
3828 f["type"] = json!(["string", lts_ms]);
3829 }
3830 "union_bool_or_string" => {
3831 f["type"] = json!(["string", "boolean"]);
3832 }
3833 _ => {}
3834 }
3835 }
3836 let reader_schema = AvroSchema::new(root.to_string());
3837 let resolved = read_alltypes_with_reader_schema(path, reader_schema);
3838
3839 fn branch_token(dt: &DataType) -> String {
3840 match dt {
3841 DataType::Null => "null".into(),
3842 DataType::Boolean => "boolean".into(),
3843 DataType::Int32 => "int".into(),
3844 DataType::Int64 => "long".into(),
3845 DataType::Float32 => "float".into(),
3846 DataType::Float64 => "double".into(),
3847 DataType::Binary => "bytes".into(),
3848 DataType::Utf8 => "string".into(),
3849 DataType::Date32 => "date".into(),
3850 DataType::Time32(arrow_schema::TimeUnit::Millisecond) => "time-millis".into(),
3851 DataType::Time64(arrow_schema::TimeUnit::Microsecond) => "time-micros".into(),
3852 DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, tz) => if tz.is_some() {
3853 "timestamp-millis"
3854 } else {
3855 "local-timestamp-millis"
3856 }
3857 .into(),
3858 DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, tz) => if tz.is_some() {
3859 "timestamp-micros"
3860 } else {
3861 "local-timestamp-micros"
3862 }
3863 .into(),
3864 DataType::Interval(IntervalUnit::MonthDayNano) => "duration".into(),
3865 DataType::FixedSizeBinary(n) => format!("fixed{n}"),
3866 DataType::Dictionary(_, _) => "enum".into(),
3867 DataType::Decimal128(p, s) => format!("decimal({p},{s})"),
3868 DataType::Decimal256(p, s) => format!("decimal({p},{s})"),
3869 #[cfg(feature = "small_decimals")]
3870 DataType::Decimal64(p, s) => format!("decimal({p},{s})"),
3871 DataType::Struct(fields) => {
3872 if fields.len() == 2 && fields[0].name() == "a" && fields[1].name() == "b" {
3873 "record:RecA".into()
3874 } else if fields.len() == 2
3875 && fields[0].name() == "x"
3876 && fields[1].name() == "y"
3877 {
3878 "record:RecB".into()
3879 } else {
3880 "record".into()
3881 }
3882 }
3883 DataType::List(_) => "array".into(),
3884 DataType::Map(_, _) => "map".into(),
3885 other => format!("{other:?}"),
3886 }
3887 }
3888
3889 fn union_tokens(u: &UnionArray) -> (Vec<i8>, HashMap<i8, String>) {
3890 let fields = match u.data_type() {
3891 DataType::Union(fields, _) => fields,
3892 other => panic!("expected Union, got {other:?}"),
3893 };
3894 let mut dict: HashMap<i8, String> = HashMap::with_capacity(fields.len());
3895 for (tid, f) in fields.iter() {
3896 dict.insert(tid, branch_token(f.data_type()));
3897 }
3898 let ids: Vec<i8> = u.type_ids().iter().copied().collect();
3899 (ids, dict)
3900 }
3901
3902 fn expected_token(field_name: &str, writer_token: &str) -> String {
3903 match field_name {
3904 "union_prim" => match writer_token {
3905 "int" => "long".into(),
3906 "float" => "double".into(),
3907 other => other.into(),
3908 },
3909 "record_with_union_field.u" => match writer_token {
3910 "int" => "long".into(),
3911 other => other.into(),
3912 },
3913 _ => writer_token.into(),
3914 }
3915 }
3916
3917 fn get_union<'a>(
3918 rb: &'a RecordBatch,
3919 schema: arrow_schema::SchemaRef,
3920 fname: &str,
3921 ) -> &'a UnionArray {
3922 let idx = schema.index_of(fname).unwrap();
3923 rb.column(idx)
3924 .as_any()
3925 .downcast_ref::<UnionArray>()
3926 .unwrap_or_else(|| panic!("{fname} should be a Union"))
3927 }
3928
3929 fn assert_union_equivalent(field_name: &str, u_writer: &UnionArray, u_reader: &UnionArray) {
3930 let (ids_w, dict_w) = union_tokens(u_writer);
3931 let (ids_r, dict_r) = union_tokens(u_reader);
3932 assert_eq!(
3933 ids_w.len(),
3934 ids_r.len(),
3935 "{field_name}: row count mismatch between baseline and resolved"
3936 );
3937 for (i, (id_w, id_r)) in ids_w.iter().zip(ids_r.iter()).enumerate() {
3938 let w_tok = dict_w.get(id_w).unwrap();
3939 let want = expected_token(field_name, w_tok);
3940 let got = dict_r.get(id_r).unwrap();
3941 assert_eq!(
3942 got, &want,
3943 "{field_name}: row {i} resolved to wrong union branch (writer={w_tok}, expected={want}, got={got})"
3944 );
3945 }
3946 }
3947
3948 for (fname, dt) in [
3949 ("nullable_int_nullfirst", DataType::Int32),
3950 ("nullable_string_nullsecond", DataType::Utf8),
3951 ] {
3952 let idx_b = baseline_schema.index_of(fname).unwrap();
3953 let idx_r = resolved.schema().index_of(fname).unwrap();
3954 let col_b = baseline.column(idx_b);
3955 let col_r = resolved.column(idx_r);
3956 assert_eq!(
3957 col_b.data_type(),
3958 &dt,
3959 "baseline {fname} should decode as non-union with nullability"
3960 );
3961 assert_eq!(
3962 col_b.as_ref(),
3963 col_r.as_ref(),
3964 "{fname}: values must be identical regardless of null-branch order"
3965 );
3966 }
3967 let union_fields = [
3968 "union_prim",
3969 "union_bytes_vs_string",
3970 "union_fixed_dur_decfix",
3971 "union_enum_records_array_map",
3972 "union_date_or_fixed4",
3973 "union_time_millis_or_enum",
3974 "union_time_micros_or_string",
3975 "union_ts_millis_utc_or_array",
3976 "union_ts_micros_local_or_bytes",
3977 "union_uuid_or_fixed10",
3978 "union_dec_bytes_or_dec_fixed",
3979 "union_null_bytes_string",
3980 "union_ts_micros_utc_or_map",
3981 "union_ts_millis_local_or_string",
3982 "union_bool_or_string",
3983 ];
3984 for fname in union_fields {
3985 let u_b = get_union(&baseline, baseline_schema.clone(), fname);
3986 let u_r = get_union(&resolved, resolved.schema(), fname);
3987 assert_union_equivalent(fname, u_b, u_r);
3988 }
3989 {
3990 let fname = "array_of_union";
3991 let idx_b = baseline_schema.index_of(fname).unwrap();
3992 let idx_r = resolved.schema().index_of(fname).unwrap();
3993 let arr_b = baseline
3994 .column(idx_b)
3995 .as_any()
3996 .downcast_ref::<ListArray>()
3997 .expect("array_of_union should be a List");
3998 let arr_r = resolved
3999 .column(idx_r)
4000 .as_any()
4001 .downcast_ref::<ListArray>()
4002 .expect("array_of_union should be a List");
4003 assert_eq!(
4004 arr_b.value_offsets(),
4005 arr_r.value_offsets(),
4006 "{fname}: list offsets changed after resolution"
4007 );
4008 let u_b = arr_b
4009 .values()
4010 .as_any()
4011 .downcast_ref::<UnionArray>()
4012 .expect("array items should be Union");
4013 let u_r = arr_r
4014 .values()
4015 .as_any()
4016 .downcast_ref::<UnionArray>()
4017 .expect("array items should be Union");
4018 let (ids_b, dict_b) = union_tokens(u_b);
4019 let (ids_r, dict_r) = union_tokens(u_r);
4020 assert_eq!(ids_b.len(), ids_r.len(), "{fname}: values length mismatch");
4021 for (i, (id_b, id_r)) in ids_b.iter().zip(ids_r.iter()).enumerate() {
4022 let w_tok = dict_b.get(id_b).unwrap();
4023 let got = dict_r.get(id_r).unwrap();
4024 assert_eq!(
4025 got, w_tok,
4026 "{fname}: value {i} resolved to wrong branch (writer={w_tok}, got={got})"
4027 );
4028 }
4029 }
4030 {
4031 let fname = "map_of_union";
4032 let idx_b = baseline_schema.index_of(fname).unwrap();
4033 let idx_r = resolved.schema().index_of(fname).unwrap();
4034 let map_b = baseline
4035 .column(idx_b)
4036 .as_any()
4037 .downcast_ref::<MapArray>()
4038 .expect("map_of_union should be a Map");
4039 let map_r = resolved
4040 .column(idx_r)
4041 .as_any()
4042 .downcast_ref::<MapArray>()
4043 .expect("map_of_union should be a Map");
4044 assert_eq!(
4045 map_b.value_offsets(),
4046 map_r.value_offsets(),
4047 "{fname}: map value offsets changed after resolution"
4048 );
4049 let ent_b = map_b.entries();
4050 let ent_r = map_r.entries();
4051 let val_b_any = ent_b.column(1).as_ref();
4052 let val_r_any = ent_r.column(1).as_ref();
4053 let b_union = val_b_any.as_any().downcast_ref::<UnionArray>();
4054 let r_union = val_r_any.as_any().downcast_ref::<UnionArray>();
4055 if let (Some(u_b), Some(u_r)) = (b_union, r_union) {
4056 assert_union_equivalent(fname, u_b, u_r);
4057 } else {
4058 assert_eq!(
4059 val_b_any.data_type(),
4060 val_r_any.data_type(),
4061 "{fname}: value data types differ after resolution"
4062 );
4063 assert_eq!(
4064 val_b_any, val_r_any,
4065 "{fname}: value arrays differ after resolution (nullable value column case)"
4066 );
4067 let value_nullable = |m: &MapArray| -> bool {
4068 match m.data_type() {
4069 DataType::Map(entries_field, _sorted) => match entries_field.data_type() {
4070 DataType::Struct(fields) => {
4071 assert_eq!(fields.len(), 2, "entries struct must have 2 fields");
4072 assert_eq!(fields[0].name(), "key");
4073 assert_eq!(fields[1].name(), "value");
4074 fields[1].is_nullable()
4075 }
4076 other => panic!("Map entries field must be Struct, got {other:?}"),
4077 },
4078 other => panic!("expected Map data type, got {other:?}"),
4079 }
4080 };
4081 assert!(
4082 value_nullable(map_b),
4083 "{fname}: baseline Map value field should be nullable per Arrow spec"
4084 );
4085 assert!(
4086 value_nullable(map_r),
4087 "{fname}: resolved Map value field should be nullable per Arrow spec"
4088 );
4089 }
4090 }
4091 {
4092 let fname = "record_with_union_field";
4093 let idx_b = baseline_schema.index_of(fname).unwrap();
4094 let idx_r = resolved.schema().index_of(fname).unwrap();
4095 let rec_b = baseline
4096 .column(idx_b)
4097 .as_any()
4098 .downcast_ref::<StructArray>()
4099 .expect("record_with_union_field should be a Struct");
4100 let rec_r = resolved
4101 .column(idx_r)
4102 .as_any()
4103 .downcast_ref::<StructArray>()
4104 .expect("record_with_union_field should be a Struct");
4105 let u_b = rec_b
4106 .column_by_name("u")
4107 .unwrap()
4108 .as_any()
4109 .downcast_ref::<UnionArray>()
4110 .expect("field 'u' should be Union (baseline)");
4111 let u_r = rec_r
4112 .column_by_name("u")
4113 .unwrap()
4114 .as_any()
4115 .downcast_ref::<UnionArray>()
4116 .expect("field 'u' should be Union (resolved)");
4117 assert_union_equivalent("record_with_union_field.u", u_b, u_r);
4118 }
4119 }
4120
4121 #[test]
4122 fn test_union_fields_end_to_end_expected_arrays() {
4123 fn tid_by_name(fields: &UnionFields, want: &str) -> i8 {
4124 for (tid, f) in fields.iter() {
4125 if f.name() == want {
4126 return tid;
4127 }
4128 }
4129 panic!("union child '{want}' not found")
4130 }
4131
4132 fn tid_by_dt(fields: &UnionFields, pred: impl Fn(&DataType) -> bool) -> i8 {
4133 for (tid, f) in fields.iter() {
4134 if pred(f.data_type()) {
4135 return tid;
4136 }
4137 }
4138 panic!("no union child matches predicate");
4139 }
4140
4141 fn uuid16_from_str(s: &str) -> [u8; 16] {
4142 fn hex(b: u8) -> u8 {
4143 match b {
4144 b'0'..=b'9' => b - b'0',
4145 b'a'..=b'f' => b - b'a' + 10,
4146 b'A'..=b'F' => b - b'A' + 10,
4147 _ => panic!("invalid hex"),
4148 }
4149 }
4150 let mut out = [0u8; 16];
4151 let bytes = s.as_bytes();
4152 let (mut i, mut j) = (0, 0);
4153 while i < bytes.len() {
4154 if bytes[i] == b'-' {
4155 i += 1;
4156 continue;
4157 }
4158 let hi = hex(bytes[i]);
4159 let lo = hex(bytes[i + 1]);
4160 out[j] = (hi << 4) | lo;
4161 j += 1;
4162 i += 2;
4163 }
4164 assert_eq!(j, 16, "uuid must decode to 16 bytes");
4165 out
4166 }
4167
4168 fn empty_child_for(dt: &DataType) -> Arc<dyn Array> {
4169 match dt {
4170 DataType::Null => Arc::new(NullArray::new(0)),
4171 DataType::Boolean => Arc::new(BooleanArray::from(Vec::<bool>::new())),
4172 DataType::Int32 => Arc::new(Int32Array::from(Vec::<i32>::new())),
4173 DataType::Int64 => Arc::new(Int64Array::from(Vec::<i64>::new())),
4174 DataType::Float32 => Arc::new(arrow_array::Float32Array::from(Vec::<f32>::new())),
4175 DataType::Float64 => Arc::new(arrow_array::Float64Array::from(Vec::<f64>::new())),
4176 DataType::Binary => Arc::new(BinaryArray::from(Vec::<&[u8]>::new())),
4177 DataType::Utf8 => Arc::new(StringArray::from(Vec::<&str>::new())),
4178 DataType::Date32 => Arc::new(arrow_array::Date32Array::from(Vec::<i32>::new())),
4179 DataType::Time32(arrow_schema::TimeUnit::Millisecond) => {
4180 Arc::new(Time32MillisecondArray::from(Vec::<i32>::new()))
4181 }
4182 DataType::Time64(arrow_schema::TimeUnit::Microsecond) => {
4183 Arc::new(Time64MicrosecondArray::from(Vec::<i64>::new()))
4184 }
4185 DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, tz) => {
4186 let a = TimestampMillisecondArray::from(Vec::<i64>::new());
4187 Arc::new(if let Some(tz) = tz {
4188 a.with_timezone(tz.clone())
4189 } else {
4190 a
4191 })
4192 }
4193 DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, tz) => {
4194 let a = TimestampMicrosecondArray::from(Vec::<i64>::new());
4195 Arc::new(if let Some(tz) = tz {
4196 a.with_timezone(tz.clone())
4197 } else {
4198 a
4199 })
4200 }
4201 DataType::Interval(IntervalUnit::MonthDayNano) => {
4202 Arc::new(arrow_array::IntervalMonthDayNanoArray::from(Vec::<
4203 IntervalMonthDayNano,
4204 >::new(
4205 )))
4206 }
4207 DataType::FixedSizeBinary(n) => Arc::new(FixedSizeBinaryArray::new_null(*n, 0)),
4208 DataType::Dictionary(k, v) => {
4209 assert_eq!(**k, DataType::Int32, "expect int32 keys for enums");
4210 let keys = Int32Array::from(Vec::<i32>::new());
4211 let values = match v.as_ref() {
4212 DataType::Utf8 => {
4213 Arc::new(StringArray::from(Vec::<&str>::new())) as ArrayRef
4214 }
4215 other => panic!("unexpected dictionary value type {other:?}"),
4216 };
4217 Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap())
4218 }
4219 DataType::List(field) => {
4220 let values: ArrayRef = match field.data_type() {
4221 DataType::Int32 => {
4222 Arc::new(Int32Array::from(Vec::<i32>::new())) as ArrayRef
4223 }
4224 DataType::Int64 => {
4225 Arc::new(Int64Array::from(Vec::<i64>::new())) as ArrayRef
4226 }
4227 DataType::Utf8 => {
4228 Arc::new(StringArray::from(Vec::<&str>::new())) as ArrayRef
4229 }
4230 DataType::Union(_, _) => {
4231 let (uf, _) = if let DataType::Union(f, m) = field.data_type() {
4232 (f.clone(), m)
4233 } else {
4234 unreachable!()
4235 };
4236 let children: Vec<ArrayRef> = uf
4237 .iter()
4238 .map(|(_, f)| empty_child_for(f.data_type()))
4239 .collect();
4240 Arc::new(
4241 UnionArray::try_new(
4242 uf.clone(),
4243 ScalarBuffer::<i8>::from(Vec::<i8>::new()),
4244 Some(ScalarBuffer::<i32>::from(Vec::<i32>::new())),
4245 children,
4246 )
4247 .unwrap(),
4248 ) as ArrayRef
4249 }
4250 other => panic!("unsupported list item type: {other:?}"),
4251 };
4252 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0]));
4253 Arc::new(ListArray::try_new(field.clone(), offsets, values, None).unwrap())
4254 }
4255 DataType::Map(entry_field, ordered) => {
4256 let DataType::Struct(childs) = entry_field.data_type() else {
4257 panic!("map entries must be struct")
4258 };
4259 let key_field = &childs[0];
4260 let val_field = &childs[1];
4261 assert_eq!(key_field.data_type(), &DataType::Utf8);
4262 let keys = StringArray::from(Vec::<&str>::new());
4263 let vals: ArrayRef = match val_field.data_type() {
4264 DataType::Float64 => {
4265 Arc::new(arrow_array::Float64Array::from(Vec::<f64>::new())) as ArrayRef
4266 }
4267 DataType::Int64 => {
4268 Arc::new(Int64Array::from(Vec::<i64>::new())) as ArrayRef
4269 }
4270 DataType::Utf8 => {
4271 Arc::new(StringArray::from(Vec::<&str>::new())) as ArrayRef
4272 }
4273 DataType::Union(uf, _) => {
4274 let ch: Vec<ArrayRef> = uf
4275 .iter()
4276 .map(|(_, f)| empty_child_for(f.data_type()))
4277 .collect();
4278 Arc::new(
4279 UnionArray::try_new(
4280 uf.clone(),
4281 ScalarBuffer::<i8>::from(Vec::<i8>::new()),
4282 Some(ScalarBuffer::<i32>::from(Vec::<i32>::new())),
4283 ch,
4284 )
4285 .unwrap(),
4286 ) as ArrayRef
4287 }
4288 other => panic!("unsupported map value type: {other:?}"),
4289 };
4290 let entries = StructArray::new(
4291 Fields::from(vec![key_field.as_ref().clone(), val_field.as_ref().clone()]),
4292 vec![Arc::new(keys) as ArrayRef, vals],
4293 None,
4294 );
4295 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0]));
4296 Arc::new(MapArray::new(
4297 entry_field.clone(),
4298 offsets,
4299 entries,
4300 None,
4301 *ordered,
4302 ))
4303 }
4304 other => panic!("empty_child_for: unhandled type {other:?}"),
4305 }
4306 }
4307
4308 fn mk_dense_union(
4309 fields: &UnionFields,
4310 type_ids: Vec<i8>,
4311 offsets: Vec<i32>,
4312 provide: impl Fn(&Field) -> Option<ArrayRef>,
4313 ) -> ArrayRef {
4314 let children: Vec<ArrayRef> = fields
4315 .iter()
4316 .map(|(_, f)| provide(f).unwrap_or_else(|| empty_child_for(f.data_type())))
4317 .collect();
4318
4319 Arc::new(
4320 UnionArray::try_new(
4321 fields.clone(),
4322 ScalarBuffer::<i8>::from(type_ids),
4323 Some(ScalarBuffer::<i32>::from(offsets)),
4324 children,
4325 )
4326 .unwrap(),
4327 ) as ArrayRef
4328 }
4329
4330 let date_a: i32 = 19_000;
4332 let time_ms_a: i32 = 13 * 3_600_000 + 45 * 60_000 + 30_000 + 123;
4333 let time_us_b: i64 = 23 * 3_600_000_000 + 59 * 60_000_000 + 59 * 1_000_000 + 999_999;
4334 let ts_ms_2024_01_01: i64 = 1_704_067_200_000;
4335 let ts_us_2024_01_01: i64 = ts_ms_2024_01_01 * 1000;
4336 let fx8_a: [u8; 8] = *b"ABCDEFGH";
4338 let fx4_abcd: [u8; 4] = *b"ABCD";
4339 let fx4_misc: [u8; 4] = [0x00, 0x11, 0x22, 0x33];
4340 let fx10_ascii: [u8; 10] = *b"0123456789";
4341 let fx10_aa: [u8; 10] = [0xAA; 10];
4342 let dur_a = IntervalMonthDayNanoType::make_value(1, 2, 3_000_000_000);
4344 let dur_b = IntervalMonthDayNanoType::make_value(12, 31, 999_000_000);
4345 let uuid1 = uuid16_from_str("fe7bc30b-4ce8-4c5e-b67c-2234a2d38e66");
4347 let uuid2 = uuid16_from_str("0826cc06-d2e3-4599-b4ad-af5fa6905cdb");
4348 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";
4354 let actual = read_file(path, 1024, false);
4355 let schema = actual.schema();
4356 let get_union = |name: &str| -> (UnionFields, UnionMode) {
4358 let idx = schema.index_of(name).unwrap();
4359 match schema.field(idx).data_type() {
4360 DataType::Union(f, m) => (f.clone(), *m),
4361 other => panic!("{name} should be a Union, got {other:?}"),
4362 }
4363 };
4364 let mut expected_cols: Vec<ArrayRef> = Vec::with_capacity(schema.fields().len());
4365 expected_cols.push(Arc::new(Int32Array::from(vec![
4367 None,
4368 Some(42),
4369 None,
4370 Some(0),
4371 ])));
4372 expected_cols.push(Arc::new(StringArray::from(vec![
4374 Some("s1"),
4375 None,
4376 Some("s3"),
4377 Some(""),
4378 ])));
4379 {
4381 let (uf, mode) = get_union("union_prim");
4382 assert!(matches!(mode, UnionMode::Dense));
4383 let generated_names: Vec<&str> = uf.iter().map(|(_, f)| f.name().as_str()).collect();
4384 let expected_names = vec![
4385 "boolean", "int", "long", "float", "double", "bytes", "string",
4386 ];
4387 assert_eq!(
4388 generated_names, expected_names,
4389 "Field names for union_prim are incorrect"
4390 );
4391 let tids = vec![
4392 tid_by_name(&uf, "long"),
4393 tid_by_name(&uf, "int"),
4394 tid_by_name(&uf, "float"),
4395 tid_by_name(&uf, "double"),
4396 ];
4397 let offs = vec![0, 0, 0, 0];
4398 let arr = mk_dense_union(&uf, tids, offs, |f| match f.name().as_str() {
4399 "int" => Some(Arc::new(Int32Array::from(vec![-1])) as ArrayRef),
4400 "long" => Some(Arc::new(Int64Array::from(vec![1_234_567_890_123i64])) as ArrayRef),
4401 "float" => {
4402 Some(Arc::new(arrow_array::Float32Array::from(vec![1.25f32])) as ArrayRef)
4403 }
4404 "double" => {
4405 Some(Arc::new(arrow_array::Float64Array::from(vec![-2.5f64])) as ArrayRef)
4406 }
4407 _ => None,
4408 });
4409 expected_cols.push(arr);
4410 }
4411 {
4413 let (uf, _) = get_union("union_bytes_vs_string");
4414 let tids = vec![
4415 tid_by_name(&uf, "bytes"),
4416 tid_by_name(&uf, "string"),
4417 tid_by_name(&uf, "string"),
4418 tid_by_name(&uf, "bytes"),
4419 ];
4420 let offs = vec![0, 0, 1, 1];
4421 let arr = mk_dense_union(&uf, tids, offs, |f| match f.name().as_str() {
4422 "bytes" => Some(
4423 Arc::new(BinaryArray::from(vec![&[0x00, 0xFF, 0x7F][..], &[][..]])) as ArrayRef,
4424 ),
4425 "string" => Some(Arc::new(StringArray::from(vec!["hello", "world"])) as ArrayRef),
4426 _ => None,
4427 });
4428 expected_cols.push(arr);
4429 }
4430 {
4432 let (uf, _) = get_union("union_fixed_dur_decfix");
4433 let tid_fx8 = tid_by_dt(&uf, |dt| matches!(dt, DataType::FixedSizeBinary(8)));
4434 let tid_dur = tid_by_dt(&uf, |dt| {
4435 matches!(
4436 dt,
4437 DataType::Interval(arrow_schema::IntervalUnit::MonthDayNano)
4438 )
4439 });
4440 let tid_dec = tid_by_dt(&uf, |dt| match dt {
4441 #[cfg(feature = "small_decimals")]
4442 DataType::Decimal64(10, 2) => true,
4443 DataType::Decimal128(10, 2) | DataType::Decimal256(10, 2) => true,
4444 _ => false,
4445 });
4446 let tids = vec![tid_fx8, tid_dur, tid_dec, tid_dur];
4447 let offs = vec![0, 0, 0, 1];
4448 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
4449 DataType::FixedSizeBinary(8) => {
4450 let it = [Some(fx8_a)].into_iter();
4451 Some(Arc::new(
4452 FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 8).unwrap(),
4453 ) as ArrayRef)
4454 }
4455 DataType::Interval(IntervalUnit::MonthDayNano) => {
4456 Some(Arc::new(arrow_array::IntervalMonthDayNanoArray::from(vec![
4457 dur_a, dur_b,
4458 ])) as ArrayRef)
4459 }
4460 #[cfg(feature = "small_decimals")]
4461 DataType::Decimal64(10, 2) => {
4462 let a = arrow_array::Decimal64Array::from_iter_values([dec_fix16_neg as i64]);
4463 Some(Arc::new(a.with_precision_and_scale(10, 2).unwrap()) as ArrayRef)
4464 }
4465 DataType::Decimal128(10, 2) => {
4466 let a = arrow_array::Decimal128Array::from_iter_values([dec_fix16_neg]);
4467 Some(Arc::new(a.with_precision_and_scale(10, 2).unwrap()) as ArrayRef)
4468 }
4469 DataType::Decimal256(10, 2) => {
4470 let a = arrow_array::Decimal256Array::from_iter_values([i256::from_i128(
4471 dec_fix16_neg,
4472 )]);
4473 Some(Arc::new(a.with_precision_and_scale(10, 2).unwrap()) as ArrayRef)
4474 }
4475 _ => None,
4476 });
4477 let generated_names: Vec<&str> = uf.iter().map(|(_, f)| f.name().as_str()).collect();
4478 let expected_names = vec!["Fx8", "Dur12", "DecFix16"];
4479 assert_eq!(
4480 generated_names, expected_names,
4481 "Data type names were not generated correctly for union_fixed_dur_decfix"
4482 );
4483 expected_cols.push(arr);
4484 }
4485 {
4487 let (uf, _) = get_union("union_enum_records_array_map");
4488 let tid_enum = tid_by_dt(&uf, |dt| matches!(dt, DataType::Dictionary(_, _)));
4489 let tid_reca = tid_by_dt(&uf, |dt| {
4490 if let DataType::Struct(fs) = dt {
4491 fs.len() == 2 && fs[0].name() == "a" && fs[1].name() == "b"
4492 } else {
4493 false
4494 }
4495 });
4496 let tid_recb = tid_by_dt(&uf, |dt| {
4497 if let DataType::Struct(fs) = dt {
4498 fs.len() == 2 && fs[0].name() == "x" && fs[1].name() == "y"
4499 } else {
4500 false
4501 }
4502 });
4503 let tid_arr = tid_by_dt(&uf, |dt| matches!(dt, DataType::List(_)));
4504 let tids = vec![tid_enum, tid_reca, tid_recb, tid_arr];
4505 let offs = vec![0, 0, 0, 0];
4506 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
4507 DataType::Dictionary(_, _) => {
4508 let keys = Int32Array::from(vec![0i32]); let values =
4510 Arc::new(StringArray::from(vec!["RED", "GREEN", "BLUE"])) as ArrayRef;
4511 Some(
4512 Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap())
4513 as ArrayRef,
4514 )
4515 }
4516 DataType::Struct(fs)
4517 if fs.len() == 2 && fs[0].name() == "a" && fs[1].name() == "b" =>
4518 {
4519 let a = Int32Array::from(vec![7]);
4520 let b = StringArray::from(vec!["x"]);
4521 Some(Arc::new(StructArray::new(
4522 fs.clone(),
4523 vec![Arc::new(a), Arc::new(b)],
4524 None,
4525 )) as ArrayRef)
4526 }
4527 DataType::Struct(fs)
4528 if fs.len() == 2 && fs[0].name() == "x" && fs[1].name() == "y" =>
4529 {
4530 let x = Int64Array::from(vec![123_456_789i64]);
4531 let y = BinaryArray::from(vec![&[0xFF, 0x00][..]]);
4532 Some(Arc::new(StructArray::new(
4533 fs.clone(),
4534 vec![Arc::new(x), Arc::new(y)],
4535 None,
4536 )) as ArrayRef)
4537 }
4538 DataType::List(field) => {
4539 let values = Int64Array::from(vec![1i64, 2, 3]);
4540 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3]));
4541 Some(Arc::new(
4542 ListArray::try_new(field.clone(), offsets, Arc::new(values), None).unwrap(),
4543 ) as ArrayRef)
4544 }
4545 DataType::Map(_, _) => None,
4546 other => panic!("unexpected child {other:?}"),
4547 });
4548 expected_cols.push(arr);
4549 }
4550 {
4552 let (uf, _) = get_union("union_date_or_fixed4");
4553 let tid_date = tid_by_dt(&uf, |dt| matches!(dt, DataType::Date32));
4554 let tid_fx4 = tid_by_dt(&uf, |dt| matches!(dt, DataType::FixedSizeBinary(4)));
4555 let tids = vec![tid_date, tid_fx4, tid_date, tid_fx4];
4556 let offs = vec![0, 0, 1, 1];
4557 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
4558 DataType::Date32 => {
4559 Some(Arc::new(arrow_array::Date32Array::from(vec![date_a, 0])) as ArrayRef)
4560 }
4561 DataType::FixedSizeBinary(4) => {
4562 let it = [Some(fx4_abcd), Some(fx4_misc)].into_iter();
4563 Some(Arc::new(
4564 FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 4).unwrap(),
4565 ) as ArrayRef)
4566 }
4567 _ => None,
4568 });
4569 expected_cols.push(arr);
4570 }
4571 {
4573 let (uf, _) = get_union("union_time_millis_or_enum");
4574 let tid_ms = tid_by_dt(&uf, |dt| {
4575 matches!(dt, DataType::Time32(arrow_schema::TimeUnit::Millisecond))
4576 });
4577 let tid_en = tid_by_dt(&uf, |dt| matches!(dt, DataType::Dictionary(_, _)));
4578 let tids = vec![tid_ms, tid_en, tid_en, tid_ms];
4579 let offs = vec![0, 0, 1, 1];
4580 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
4581 DataType::Time32(arrow_schema::TimeUnit::Millisecond) => {
4582 Some(Arc::new(Time32MillisecondArray::from(vec![time_ms_a, 0])) as ArrayRef)
4583 }
4584 DataType::Dictionary(_, _) => {
4585 let keys = Int32Array::from(vec![0i32, 1]); let values = Arc::new(StringArray::from(vec!["ON", "OFF"])) as ArrayRef;
4587 Some(
4588 Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap())
4589 as ArrayRef,
4590 )
4591 }
4592 _ => None,
4593 });
4594 expected_cols.push(arr);
4595 }
4596 {
4598 let (uf, _) = get_union("union_time_micros_or_string");
4599 let tid_us = tid_by_dt(&uf, |dt| {
4600 matches!(dt, DataType::Time64(arrow_schema::TimeUnit::Microsecond))
4601 });
4602 let tid_s = tid_by_name(&uf, "string");
4603 let tids = vec![tid_s, tid_us, tid_s, tid_s];
4604 let offs = vec![0, 0, 1, 2];
4605 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
4606 DataType::Time64(arrow_schema::TimeUnit::Microsecond) => {
4607 Some(Arc::new(Time64MicrosecondArray::from(vec![time_us_b])) as ArrayRef)
4608 }
4609 DataType::Utf8 => {
4610 Some(Arc::new(StringArray::from(vec!["evening", "night", ""])) as ArrayRef)
4611 }
4612 _ => None,
4613 });
4614 expected_cols.push(arr);
4615 }
4616 {
4618 let (uf, _) = get_union("union_ts_millis_utc_or_array");
4619 let tid_ts = tid_by_dt(&uf, |dt| {
4620 matches!(
4621 dt,
4622 DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, _)
4623 )
4624 });
4625 let tid_arr = tid_by_dt(&uf, |dt| matches!(dt, DataType::List(_)));
4626 let tids = vec![tid_ts, tid_arr, tid_arr, tid_ts];
4627 let offs = vec![0, 0, 1, 1];
4628 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
4629 DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, tz) => {
4630 let a = TimestampMillisecondArray::from(vec![
4631 ts_ms_2024_01_01,
4632 ts_ms_2024_01_01 + 86_400_000,
4633 ]);
4634 Some(Arc::new(if let Some(tz) = tz {
4635 a.with_timezone(tz.clone())
4636 } else {
4637 a
4638 }) as ArrayRef)
4639 }
4640 DataType::List(field) => {
4641 let values = Int32Array::from(vec![0, 1, 2, -1, 0, 1]);
4642 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 6]));
4643 Some(Arc::new(
4644 ListArray::try_new(field.clone(), offsets, Arc::new(values), None).unwrap(),
4645 ) as ArrayRef)
4646 }
4647 _ => None,
4648 });
4649 expected_cols.push(arr);
4650 }
4651 {
4653 let (uf, _) = get_union("union_ts_micros_local_or_bytes");
4654 let tid_lts = tid_by_dt(&uf, |dt| {
4655 matches!(
4656 dt,
4657 DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, None)
4658 )
4659 });
4660 let tid_b = tid_by_name(&uf, "bytes");
4661 let tids = vec![tid_b, tid_lts, tid_b, tid_b];
4662 let offs = vec![0, 0, 1, 2];
4663 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
4664 DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, None) => Some(Arc::new(
4665 TimestampMicrosecondArray::from(vec![ts_us_2024_01_01]),
4666 )
4667 as ArrayRef),
4668 DataType::Binary => Some(Arc::new(BinaryArray::from(vec![
4669 &b"\x11\x22\x33"[..],
4670 &b"\x00"[..],
4671 &b"\x10\x20\x30\x40"[..],
4672 ])) as ArrayRef),
4673 _ => None,
4674 });
4675 expected_cols.push(arr);
4676 }
4677 {
4679 let (uf, _) = get_union("union_uuid_or_fixed10");
4680 let tid_fx16 = tid_by_dt(&uf, |dt| matches!(dt, DataType::FixedSizeBinary(16)));
4681 let tid_fx10 = tid_by_dt(&uf, |dt| matches!(dt, DataType::FixedSizeBinary(10)));
4682 let tids = vec![tid_fx16, tid_fx10, tid_fx16, tid_fx10];
4683 let offs = vec![0, 0, 1, 1];
4684 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
4685 DataType::FixedSizeBinary(16) => {
4686 let it = [Some(uuid1), Some(uuid2)].into_iter();
4687 Some(Arc::new(
4688 FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 16).unwrap(),
4689 ) as ArrayRef)
4690 }
4691 DataType::FixedSizeBinary(10) => {
4692 let it = [Some(fx10_ascii), Some(fx10_aa)].into_iter();
4693 Some(Arc::new(
4694 FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 10).unwrap(),
4695 ) as ArrayRef)
4696 }
4697 _ => None,
4698 });
4699 expected_cols.push(arr);
4700 }
4701 {
4703 let (uf, _) = get_union("union_dec_bytes_or_dec_fixed");
4704 let tid_b10s2 = tid_by_dt(&uf, |dt| match dt {
4705 #[cfg(feature = "small_decimals")]
4706 DataType::Decimal64(10, 2) => true,
4707 DataType::Decimal128(10, 2) | DataType::Decimal256(10, 2) => true,
4708 _ => false,
4709 });
4710 let tid_f20s4 = tid_by_dt(&uf, |dt| {
4711 matches!(
4712 dt,
4713 DataType::Decimal128(20, 4) | DataType::Decimal256(20, 4)
4714 )
4715 });
4716 let tids = vec![tid_b10s2, tid_f20s4, tid_b10s2, tid_f20s4];
4717 let offs = vec![0, 0, 1, 1];
4718 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
4719 #[cfg(feature = "small_decimals")]
4720 DataType::Decimal64(10, 2) => {
4721 let a = Decimal64Array::from_iter_values([dec_b_scale2_pos as i64, 0i64]);
4722 Some(Arc::new(a.with_precision_and_scale(10, 2).unwrap()) as ArrayRef)
4723 }
4724 DataType::Decimal128(10, 2) => {
4725 let a = Decimal128Array::from_iter_values([dec_b_scale2_pos, 0]);
4726 Some(Arc::new(a.with_precision_and_scale(10, 2).unwrap()) as ArrayRef)
4727 }
4728 DataType::Decimal256(10, 2) => {
4729 let a = Decimal256Array::from_iter_values([
4730 i256::from_i128(dec_b_scale2_pos),
4731 i256::from(0),
4732 ]);
4733 Some(Arc::new(a.with_precision_and_scale(10, 2).unwrap()) as ArrayRef)
4734 }
4735 DataType::Decimal128(20, 4) => {
4736 let a = Decimal128Array::from_iter_values([dec_fix20_s4_neg, dec_fix20_s4]);
4737 Some(Arc::new(a.with_precision_and_scale(20, 4).unwrap()) as ArrayRef)
4738 }
4739 DataType::Decimal256(20, 4) => {
4740 let a = Decimal256Array::from_iter_values([
4741 i256::from_i128(dec_fix20_s4_neg),
4742 i256::from_i128(dec_fix20_s4),
4743 ]);
4744 Some(Arc::new(a.with_precision_and_scale(20, 4).unwrap()) as ArrayRef)
4745 }
4746 _ => None,
4747 });
4748 expected_cols.push(arr);
4749 }
4750 {
4752 let (uf, _) = get_union("union_null_bytes_string");
4753 let tid_n = tid_by_name(&uf, "null");
4754 let tid_b = tid_by_name(&uf, "bytes");
4755 let tid_s = tid_by_name(&uf, "string");
4756 let tids = vec![tid_n, tid_b, tid_s, tid_s];
4757 let offs = vec![0, 0, 0, 1];
4758 let arr = mk_dense_union(&uf, tids, offs, |f| match f.name().as_str() {
4759 "null" => Some(Arc::new(arrow_array::NullArray::new(1)) as ArrayRef),
4760 "bytes" => Some(Arc::new(BinaryArray::from(vec![&b"\x01\x02"[..]])) as ArrayRef),
4761 "string" => Some(Arc::new(StringArray::from(vec!["text", "u"])) as ArrayRef),
4762 _ => None,
4763 });
4764 expected_cols.push(arr);
4765 }
4766 {
4768 let idx = schema.index_of("array_of_union").unwrap();
4769 let dt = schema.field(idx).data_type().clone();
4770 let (item_field, _) = match &dt {
4771 DataType::List(f) => (f.clone(), ()),
4772 other => panic!("array_of_union must be List, got {other:?}"),
4773 };
4774 let (uf, _) = match item_field.data_type() {
4775 DataType::Union(f, m) => (f.clone(), m),
4776 other => panic!("array_of_union items must be Union, got {other:?}"),
4777 };
4778 let tid_l = tid_by_name(&uf, "long");
4779 let tid_s = tid_by_name(&uf, "string");
4780 let type_ids = vec![tid_l, tid_s, tid_l, tid_s, tid_l, tid_l, tid_s, tid_l];
4781 let offsets = vec![0, 0, 1, 1, 2, 3, 2, 4];
4782 let values_union =
4783 mk_dense_union(&uf, type_ids, offsets, |f| match f.name().as_str() {
4784 "long" => {
4785 Some(Arc::new(Int64Array::from(vec![1i64, -5, 42, -1, 0])) as ArrayRef)
4786 }
4787 "string" => Some(Arc::new(StringArray::from(vec!["a", "", "z"])) as ArrayRef),
4788 _ => None,
4789 });
4790 let list_offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 5, 6, 8]));
4791 expected_cols.push(Arc::new(
4792 ListArray::try_new(item_field.clone(), list_offsets, values_union, None).unwrap(),
4793 ));
4794 }
4795 {
4797 let idx = schema.index_of("map_of_union").unwrap();
4798 let dt = schema.field(idx).data_type().clone();
4799 let (entry_field, ordered) = match &dt {
4800 DataType::Map(f, ordered) => (f.clone(), *ordered),
4801 other => panic!("map_of_union must be Map, got {other:?}"),
4802 };
4803 let DataType::Struct(entry_fields) = entry_field.data_type() else {
4804 panic!("map entries must be struct")
4805 };
4806 let key_field = entry_fields[0].clone();
4807 let val_field = entry_fields[1].clone();
4808 let keys = StringArray::from(vec!["a", "b", "x", "pi"]);
4809 let rounded_pi = (std::f64::consts::PI * 100_000.0).round() / 100_000.0;
4810 let values: ArrayRef = match val_field.data_type() {
4811 DataType::Union(uf, _) => {
4812 let tid_n = tid_by_name(uf, "null");
4813 let tid_d = tid_by_name(uf, "double");
4814 let tids = vec![tid_n, tid_d, tid_d, tid_d];
4815 let offs = vec![0, 0, 1, 2];
4816 mk_dense_union(uf, tids, offs, |f| match f.name().as_str() {
4817 "null" => Some(Arc::new(NullArray::new(1)) as ArrayRef),
4818 "double" => Some(Arc::new(arrow_array::Float64Array::from(vec![
4819 2.5f64, -0.5f64, rounded_pi,
4820 ])) as ArrayRef),
4821 _ => None,
4822 })
4823 }
4824 DataType::Float64 => Arc::new(arrow_array::Float64Array::from(vec![
4825 None,
4826 Some(2.5),
4827 Some(-0.5),
4828 Some(rounded_pi),
4829 ])),
4830 other => panic!("unexpected map value type {other:?}"),
4831 };
4832 let entries = StructArray::new(
4833 Fields::from(vec![key_field.as_ref().clone(), val_field.as_ref().clone()]),
4834 vec![Arc::new(keys) as ArrayRef, values],
4835 None,
4836 );
4837 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2, 3, 3, 4]));
4838 expected_cols.push(Arc::new(MapArray::new(
4839 entry_field,
4840 offsets,
4841 entries,
4842 None,
4843 ordered,
4844 )));
4845 }
4846 {
4848 let idx = schema.index_of("record_with_union_field").unwrap();
4849 let DataType::Struct(rec_fields) = schema.field(idx).data_type() else {
4850 panic!("record_with_union_field should be Struct")
4851 };
4852 let id = Int32Array::from(vec![1, 2, 3, 4]);
4853 let u_field = rec_fields.iter().find(|f| f.name() == "u").unwrap();
4854 let DataType::Union(uf, _) = u_field.data_type() else {
4855 panic!("u must be Union")
4856 };
4857 let tid_i = tid_by_name(uf, "int");
4858 let tid_s = tid_by_name(uf, "string");
4859 let tids = vec![tid_s, tid_i, tid_i, tid_s];
4860 let offs = vec![0, 0, 1, 1];
4861 let u = mk_dense_union(uf, tids, offs, |f| match f.name().as_str() {
4862 "int" => Some(Arc::new(Int32Array::from(vec![99, 0])) as ArrayRef),
4863 "string" => Some(Arc::new(StringArray::from(vec!["one", "four"])) as ArrayRef),
4864 _ => None,
4865 });
4866 let rec = StructArray::new(rec_fields.clone(), vec![Arc::new(id) as ArrayRef, u], None);
4867 expected_cols.push(Arc::new(rec));
4868 }
4869 {
4871 let (uf, _) = get_union("union_ts_micros_utc_or_map");
4872 let tid_ts = tid_by_dt(&uf, |dt| {
4873 matches!(
4874 dt,
4875 DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, Some(_))
4876 )
4877 });
4878 let tid_map = tid_by_dt(&uf, |dt| matches!(dt, DataType::Map(_, _)));
4879 let tids = vec![tid_ts, tid_map, tid_ts, tid_map];
4880 let offs = vec![0, 0, 1, 1];
4881 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
4882 DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, tz) => {
4883 let a = TimestampMicrosecondArray::from(vec![ts_us_2024_01_01, 0i64]);
4884 Some(Arc::new(if let Some(tz) = tz {
4885 a.with_timezone(tz.clone())
4886 } else {
4887 a
4888 }) as ArrayRef)
4889 }
4890 DataType::Map(entry_field, ordered) => {
4891 let DataType::Struct(fs) = entry_field.data_type() else {
4892 panic!("map entries must be struct")
4893 };
4894 let key_field = fs[0].clone();
4895 let val_field = fs[1].clone();
4896 assert_eq!(key_field.data_type(), &DataType::Utf8);
4897 assert_eq!(val_field.data_type(), &DataType::Int64);
4898 let keys = StringArray::from(vec!["k1", "k2", "n"]);
4899 let vals = Int64Array::from(vec![1i64, 2, 0]);
4900 let entries = StructArray::new(
4901 Fields::from(vec![key_field.as_ref().clone(), val_field.as_ref().clone()]),
4902 vec![Arc::new(keys) as ArrayRef, Arc::new(vals) as ArrayRef],
4903 None,
4904 );
4905 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2, 3]));
4906 Some(Arc::new(MapArray::new(
4907 entry_field.clone(),
4908 offsets,
4909 entries,
4910 None,
4911 *ordered,
4912 )) as ArrayRef)
4913 }
4914 _ => None,
4915 });
4916 expected_cols.push(arr);
4917 }
4918 {
4920 let (uf, _) = get_union("union_ts_millis_local_or_string");
4921 let tid_ts = tid_by_dt(&uf, |dt| {
4922 matches!(
4923 dt,
4924 DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, None)
4925 )
4926 });
4927 let tid_s = tid_by_name(&uf, "string");
4928 let tids = vec![tid_s, tid_ts, tid_s, tid_s];
4929 let offs = vec![0, 0, 1, 2];
4930 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
4931 DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, None) => Some(Arc::new(
4932 TimestampMillisecondArray::from(vec![ts_ms_2024_01_01]),
4933 )
4934 as ArrayRef),
4935 DataType::Utf8 => {
4936 Some(
4937 Arc::new(StringArray::from(vec!["local midnight", "done", ""])) as ArrayRef,
4938 )
4939 }
4940 _ => None,
4941 });
4942 expected_cols.push(arr);
4943 }
4944 {
4946 let (uf, _) = get_union("union_bool_or_string");
4947 let tid_b = tid_by_name(&uf, "boolean");
4948 let tid_s = tid_by_name(&uf, "string");
4949 let tids = vec![tid_b, tid_s, tid_b, tid_s];
4950 let offs = vec![0, 0, 1, 1];
4951 let arr = mk_dense_union(&uf, tids, offs, |f| match f.name().as_str() {
4952 "boolean" => Some(Arc::new(BooleanArray::from(vec![true, false])) as ArrayRef),
4953 "string" => Some(Arc::new(StringArray::from(vec!["no", "yes"])) as ArrayRef),
4954 _ => None,
4955 });
4956 expected_cols.push(arr);
4957 }
4958 let expected = RecordBatch::try_new(schema.clone(), expected_cols).unwrap();
4959 assert_eq!(
4960 actual, expected,
4961 "full end-to-end equality for union_fields.avro"
4962 );
4963 }
4964
4965 #[test]
4966 fn test_read_zero_byte_avro_file() {
4967 let batch = read_file("test/data/zero_byte.avro", 3, false);
4968 let schema = batch.schema();
4969 assert_eq!(schema.fields().len(), 1);
4970 let field = schema.field(0);
4971 assert_eq!(field.name(), "data");
4972 assert_eq!(field.data_type(), &DataType::Binary);
4973 assert!(field.is_nullable());
4974 assert_eq!(batch.num_rows(), 3);
4975 assert_eq!(batch.num_columns(), 1);
4976 let binary_array = batch
4977 .column(0)
4978 .as_any()
4979 .downcast_ref::<BinaryArray>()
4980 .unwrap();
4981 assert!(binary_array.is_null(0));
4982 assert!(binary_array.is_valid(1));
4983 assert_eq!(binary_array.value(1), b"");
4984 assert!(binary_array.is_valid(2));
4985 assert_eq!(binary_array.value(2), b"some bytes");
4986 }
4987
4988 #[test]
4989 fn test_alltypes() {
4990 let expected = RecordBatch::try_from_iter_with_nullable([
4991 (
4992 "id",
4993 Arc::new(Int32Array::from(vec![4, 5, 6, 7, 2, 3, 0, 1])) as _,
4994 true,
4995 ),
4996 (
4997 "bool_col",
4998 Arc::new(BooleanArray::from_iter((0..8).map(|x| Some(x % 2 == 0)))) as _,
4999 true,
5000 ),
5001 (
5002 "tinyint_col",
5003 Arc::new(Int32Array::from_iter_values((0..8).map(|x| x % 2))) as _,
5004 true,
5005 ),
5006 (
5007 "smallint_col",
5008 Arc::new(Int32Array::from_iter_values((0..8).map(|x| x % 2))) as _,
5009 true,
5010 ),
5011 (
5012 "int_col",
5013 Arc::new(Int32Array::from_iter_values((0..8).map(|x| x % 2))) as _,
5014 true,
5015 ),
5016 (
5017 "bigint_col",
5018 Arc::new(Int64Array::from_iter_values((0..8).map(|x| (x % 2) * 10))) as _,
5019 true,
5020 ),
5021 (
5022 "float_col",
5023 Arc::new(Float32Array::from_iter_values(
5024 (0..8).map(|x| (x % 2) as f32 * 1.1),
5025 )) as _,
5026 true,
5027 ),
5028 (
5029 "double_col",
5030 Arc::new(Float64Array::from_iter_values(
5031 (0..8).map(|x| (x % 2) as f64 * 10.1),
5032 )) as _,
5033 true,
5034 ),
5035 (
5036 "date_string_col",
5037 Arc::new(BinaryArray::from_iter_values([
5038 [48, 51, 47, 48, 49, 47, 48, 57],
5039 [48, 51, 47, 48, 49, 47, 48, 57],
5040 [48, 52, 47, 48, 49, 47, 48, 57],
5041 [48, 52, 47, 48, 49, 47, 48, 57],
5042 [48, 50, 47, 48, 49, 47, 48, 57],
5043 [48, 50, 47, 48, 49, 47, 48, 57],
5044 [48, 49, 47, 48, 49, 47, 48, 57],
5045 [48, 49, 47, 48, 49, 47, 48, 57],
5046 ])) as _,
5047 true,
5048 ),
5049 (
5050 "string_col",
5051 Arc::new(BinaryArray::from_iter_values((0..8).map(|x| [48 + x % 2]))) as _,
5052 true,
5053 ),
5054 (
5055 "timestamp_col",
5056 Arc::new(
5057 TimestampMicrosecondArray::from_iter_values([
5058 1235865600000000, 1235865660000000, 1238544000000000, 1238544060000000, 1233446400000000, 1233446460000000, 1230768000000000, 1230768060000000, ])
5067 .with_timezone("+00:00"),
5068 ) as _,
5069 true,
5070 ),
5071 ])
5072 .unwrap();
5073
5074 for file in files() {
5075 let file = arrow_test_data(file);
5076
5077 assert_eq!(read_file(&file, 8, false), expected);
5078 assert_eq!(read_file(&file, 3, false), expected);
5079 }
5080 }
5081
5082 #[test]
5083 #[cfg(feature = "snappy")]
5085 fn test_alltypes_dictionary() {
5086 let file = "avro/alltypes_dictionary.avro";
5087 let expected = RecordBatch::try_from_iter_with_nullable([
5088 ("id", Arc::new(Int32Array::from(vec![0, 1])) as _, true),
5089 (
5090 "bool_col",
5091 Arc::new(BooleanArray::from(vec![Some(true), Some(false)])) as _,
5092 true,
5093 ),
5094 (
5095 "tinyint_col",
5096 Arc::new(Int32Array::from(vec![0, 1])) as _,
5097 true,
5098 ),
5099 (
5100 "smallint_col",
5101 Arc::new(Int32Array::from(vec![0, 1])) as _,
5102 true,
5103 ),
5104 ("int_col", Arc::new(Int32Array::from(vec![0, 1])) as _, true),
5105 (
5106 "bigint_col",
5107 Arc::new(Int64Array::from(vec![0, 10])) as _,
5108 true,
5109 ),
5110 (
5111 "float_col",
5112 Arc::new(Float32Array::from(vec![0.0, 1.1])) as _,
5113 true,
5114 ),
5115 (
5116 "double_col",
5117 Arc::new(Float64Array::from(vec![0.0, 10.1])) as _,
5118 true,
5119 ),
5120 (
5121 "date_string_col",
5122 Arc::new(BinaryArray::from_iter_values([b"01/01/09", b"01/01/09"])) as _,
5123 true,
5124 ),
5125 (
5126 "string_col",
5127 Arc::new(BinaryArray::from_iter_values([b"0", b"1"])) as _,
5128 true,
5129 ),
5130 (
5131 "timestamp_col",
5132 Arc::new(
5133 TimestampMicrosecondArray::from_iter_values([
5134 1230768000000000, 1230768060000000, ])
5137 .with_timezone("+00:00"),
5138 ) as _,
5139 true,
5140 ),
5141 ])
5142 .unwrap();
5143 let file_path = arrow_test_data(file);
5144 let batch_large = read_file(&file_path, 8, false);
5145 assert_eq!(
5146 batch_large, expected,
5147 "Decoded RecordBatch does not match for file {file}"
5148 );
5149 let batch_small = read_file(&file_path, 3, false);
5150 assert_eq!(
5151 batch_small, expected,
5152 "Decoded RecordBatch (batch size 3) does not match for file {file}"
5153 );
5154 }
5155
5156 #[test]
5157 fn test_alltypes_nulls_plain() {
5158 let file = "avro/alltypes_nulls_plain.avro";
5159 let expected = RecordBatch::try_from_iter_with_nullable([
5160 (
5161 "string_col",
5162 Arc::new(StringArray::from(vec![None::<&str>])) as _,
5163 true,
5164 ),
5165 ("int_col", Arc::new(Int32Array::from(vec![None])) as _, true),
5166 (
5167 "bool_col",
5168 Arc::new(BooleanArray::from(vec![None])) as _,
5169 true,
5170 ),
5171 (
5172 "bigint_col",
5173 Arc::new(Int64Array::from(vec![None])) as _,
5174 true,
5175 ),
5176 (
5177 "float_col",
5178 Arc::new(Float32Array::from(vec![None])) as _,
5179 true,
5180 ),
5181 (
5182 "double_col",
5183 Arc::new(Float64Array::from(vec![None])) as _,
5184 true,
5185 ),
5186 (
5187 "bytes_col",
5188 Arc::new(BinaryArray::from(vec![None::<&[u8]>])) as _,
5189 true,
5190 ),
5191 ])
5192 .unwrap();
5193 let file_path = arrow_test_data(file);
5194 let batch_large = read_file(&file_path, 8, false);
5195 assert_eq!(
5196 batch_large, expected,
5197 "Decoded RecordBatch does not match for file {file}"
5198 );
5199 let batch_small = read_file(&file_path, 3, false);
5200 assert_eq!(
5201 batch_small, expected,
5202 "Decoded RecordBatch (batch size 3) does not match for file {file}"
5203 );
5204 }
5205
5206 #[test]
5207 #[cfg(feature = "snappy")]
5209 fn test_binary() {
5210 let file = arrow_test_data("avro/binary.avro");
5211 let batch = read_file(&file, 8, false);
5212 let expected = RecordBatch::try_from_iter_with_nullable([(
5213 "foo",
5214 Arc::new(BinaryArray::from_iter_values(vec![
5215 b"\x00" as &[u8],
5216 b"\x01" as &[u8],
5217 b"\x02" as &[u8],
5218 b"\x03" as &[u8],
5219 b"\x04" as &[u8],
5220 b"\x05" as &[u8],
5221 b"\x06" as &[u8],
5222 b"\x07" as &[u8],
5223 b"\x08" as &[u8],
5224 b"\t" as &[u8],
5225 b"\n" as &[u8],
5226 b"\x0b" as &[u8],
5227 ])) as Arc<dyn Array>,
5228 true,
5229 )])
5230 .unwrap();
5231 assert_eq!(batch, expected);
5232 }
5233
5234 #[test]
5235 #[cfg(feature = "snappy")]
5237 fn test_decimal() {
5238 #[cfg(feature = "small_decimals")]
5242 let files: [(&str, DataType, HashMap<String, String>); 8] = [
5243 (
5244 "avro/fixed_length_decimal.avro",
5245 DataType::Decimal128(25, 2),
5246 HashMap::from([
5247 (
5248 "avro.namespace".to_string(),
5249 "topLevelRecord.value".to_string(),
5250 ),
5251 ("avro.name".to_string(), "fixed".to_string()),
5252 ]),
5253 ),
5254 (
5255 "avro/fixed_length_decimal_legacy.avro",
5256 DataType::Decimal64(13, 2),
5257 HashMap::from([
5258 (
5259 "avro.namespace".to_string(),
5260 "topLevelRecord.value".to_string(),
5261 ),
5262 ("avro.name".to_string(), "fixed".to_string()),
5263 ]),
5264 ),
5265 (
5266 "avro/int32_decimal.avro",
5267 DataType::Decimal32(4, 2),
5268 HashMap::from([
5269 (
5270 "avro.namespace".to_string(),
5271 "topLevelRecord.value".to_string(),
5272 ),
5273 ("avro.name".to_string(), "fixed".to_string()),
5274 ]),
5275 ),
5276 (
5277 "avro/int64_decimal.avro",
5278 DataType::Decimal64(10, 2),
5279 HashMap::from([
5280 (
5281 "avro.namespace".to_string(),
5282 "topLevelRecord.value".to_string(),
5283 ),
5284 ("avro.name".to_string(), "fixed".to_string()),
5285 ]),
5286 ),
5287 (
5288 "test/data/int256_decimal.avro",
5289 DataType::Decimal256(76, 10),
5290 HashMap::new(),
5291 ),
5292 (
5293 "test/data/fixed256_decimal.avro",
5294 DataType::Decimal256(76, 10),
5295 HashMap::from([("avro.name".to_string(), "Decimal256Fixed".to_string())]),
5296 ),
5297 (
5298 "test/data/fixed_length_decimal_legacy_32.avro",
5299 DataType::Decimal32(9, 2),
5300 HashMap::from([("avro.name".to_string(), "Decimal32FixedLegacy".to_string())]),
5301 ),
5302 (
5303 "test/data/int128_decimal.avro",
5304 DataType::Decimal128(38, 2),
5305 HashMap::new(),
5306 ),
5307 ];
5308 #[cfg(not(feature = "small_decimals"))]
5309 let files: [(&str, DataType, HashMap<String, String>); 8] = [
5310 (
5311 "avro/fixed_length_decimal.avro",
5312 DataType::Decimal128(25, 2),
5313 HashMap::from([
5314 (
5315 "avro.namespace".to_string(),
5316 "topLevelRecord.value".to_string(),
5317 ),
5318 ("avro.name".to_string(), "fixed".to_string()),
5319 ]),
5320 ),
5321 (
5322 "avro/fixed_length_decimal_legacy.avro",
5323 DataType::Decimal128(13, 2),
5324 HashMap::from([
5325 (
5326 "avro.namespace".to_string(),
5327 "topLevelRecord.value".to_string(),
5328 ),
5329 ("avro.name".to_string(), "fixed".to_string()),
5330 ]),
5331 ),
5332 (
5333 "avro/int32_decimal.avro",
5334 DataType::Decimal128(4, 2),
5335 HashMap::from([
5336 (
5337 "avro.namespace".to_string(),
5338 "topLevelRecord.value".to_string(),
5339 ),
5340 ("avro.name".to_string(), "fixed".to_string()),
5341 ]),
5342 ),
5343 (
5344 "avro/int64_decimal.avro",
5345 DataType::Decimal128(10, 2),
5346 HashMap::from([
5347 (
5348 "avro.namespace".to_string(),
5349 "topLevelRecord.value".to_string(),
5350 ),
5351 ("avro.name".to_string(), "fixed".to_string()),
5352 ]),
5353 ),
5354 (
5355 "test/data/int256_decimal.avro",
5356 DataType::Decimal256(76, 10),
5357 HashMap::new(),
5358 ),
5359 (
5360 "test/data/fixed256_decimal.avro",
5361 DataType::Decimal256(76, 10),
5362 HashMap::from([("avro.name".to_string(), "Decimal256Fixed".to_string())]),
5363 ),
5364 (
5365 "test/data/fixed_length_decimal_legacy_32.avro",
5366 DataType::Decimal128(9, 2),
5367 HashMap::from([("avro.name".to_string(), "Decimal32FixedLegacy".to_string())]),
5368 ),
5369 (
5370 "test/data/int128_decimal.avro",
5371 DataType::Decimal128(38, 2),
5372 HashMap::new(),
5373 ),
5374 ];
5375 for (file, expected_dt, mut metadata) in files {
5376 let (precision, scale) = match expected_dt {
5377 DataType::Decimal32(p, s)
5378 | DataType::Decimal64(p, s)
5379 | DataType::Decimal128(p, s)
5380 | DataType::Decimal256(p, s) => (p, s),
5381 _ => unreachable!("Unexpected decimal type in test inputs"),
5382 };
5383 assert!(scale >= 0, "test data uses non-negative scales only");
5384 let scale_u32 = scale as u32;
5385 let file_path: String = if file.starts_with("avro/") {
5386 arrow_test_data(file)
5387 } else {
5388 std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR"))
5389 .join(file)
5390 .to_string_lossy()
5391 .into_owned()
5392 };
5393 let pow10: i128 = 10i128.pow(scale_u32);
5394 let values_i128: Vec<i128> = (1..=24).map(|n| (n as i128) * pow10).collect();
5395 let build_expected = |dt: &DataType, values: &[i128]| -> ArrayRef {
5396 match *dt {
5397 #[cfg(feature = "small_decimals")]
5398 DataType::Decimal32(p, s) => {
5399 let it = values.iter().map(|&v| v as i32);
5400 Arc::new(
5401 Decimal32Array::from_iter_values(it)
5402 .with_precision_and_scale(p, s)
5403 .unwrap(),
5404 )
5405 }
5406 #[cfg(feature = "small_decimals")]
5407 DataType::Decimal64(p, s) => {
5408 let it = values.iter().map(|&v| v as i64);
5409 Arc::new(
5410 Decimal64Array::from_iter_values(it)
5411 .with_precision_and_scale(p, s)
5412 .unwrap(),
5413 )
5414 }
5415 DataType::Decimal128(p, s) => {
5416 let it = values.iter().copied();
5417 Arc::new(
5418 Decimal128Array::from_iter_values(it)
5419 .with_precision_and_scale(p, s)
5420 .unwrap(),
5421 )
5422 }
5423 DataType::Decimal256(p, s) => {
5424 let it = values.iter().map(|&v| i256::from_i128(v));
5425 Arc::new(
5426 Decimal256Array::from_iter_values(it)
5427 .with_precision_and_scale(p, s)
5428 .unwrap(),
5429 )
5430 }
5431 _ => unreachable!("Unexpected decimal type in test"),
5432 }
5433 };
5434 let actual_batch = read_file(&file_path, 8, false);
5435 let actual_nullable = actual_batch.schema().field(0).is_nullable();
5436 let expected_array = build_expected(&expected_dt, &values_i128);
5437 metadata.insert("precision".to_string(), precision.to_string());
5438 metadata.insert("scale".to_string(), scale.to_string());
5439 let field =
5440 Field::new("value", expected_dt.clone(), actual_nullable).with_metadata(metadata);
5441 let expected_schema = Arc::new(Schema::new(vec![field]));
5442 let expected_batch =
5443 RecordBatch::try_new(expected_schema.clone(), vec![expected_array]).unwrap();
5444 assert_eq!(
5445 actual_batch, expected_batch,
5446 "Decoded RecordBatch does not match for {file}"
5447 );
5448 let actual_batch_small = read_file(&file_path, 3, false);
5449 assert_eq!(
5450 actual_batch_small, expected_batch,
5451 "Decoded RecordBatch does not match for {file} with batch size 3"
5452 );
5453 }
5454 }
5455
5456 #[test]
5457 fn test_read_duration_logical_types_feature_toggle() -> Result<(), ArrowError> {
5458 let file_path = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR"))
5459 .join("test/data/duration_logical_types.avro")
5460 .to_string_lossy()
5461 .into_owned();
5462
5463 let actual_batch = read_file(&file_path, 4, false);
5464
5465 let expected_batch = {
5466 #[cfg(feature = "avro_custom_types")]
5467 {
5468 let schema = Arc::new(Schema::new(vec![
5469 Field::new(
5470 "duration_time_nanos",
5471 DataType::Duration(TimeUnit::Nanosecond),
5472 false,
5473 ),
5474 Field::new(
5475 "duration_time_micros",
5476 DataType::Duration(TimeUnit::Microsecond),
5477 false,
5478 ),
5479 Field::new(
5480 "duration_time_millis",
5481 DataType::Duration(TimeUnit::Millisecond),
5482 false,
5483 ),
5484 Field::new(
5485 "duration_time_seconds",
5486 DataType::Duration(TimeUnit::Second),
5487 false,
5488 ),
5489 ]));
5490
5491 let nanos = Arc::new(PrimitiveArray::<DurationNanosecondType>::from(vec![
5492 10, 20, 30, 40,
5493 ])) as ArrayRef;
5494 let micros = Arc::new(PrimitiveArray::<DurationMicrosecondType>::from(vec![
5495 100, 200, 300, 400,
5496 ])) as ArrayRef;
5497 let millis = Arc::new(PrimitiveArray::<DurationMillisecondType>::from(vec![
5498 1000, 2000, 3000, 4000,
5499 ])) as ArrayRef;
5500 let seconds = Arc::new(PrimitiveArray::<DurationSecondType>::from(vec![1, 2, 3, 4]))
5501 as ArrayRef;
5502
5503 RecordBatch::try_new(schema, vec![nanos, micros, millis, seconds])?
5504 }
5505 #[cfg(not(feature = "avro_custom_types"))]
5506 {
5507 let schema = Arc::new(Schema::new(vec![
5508 Field::new("duration_time_nanos", DataType::Int64, false)
5509 .with_metadata([("logicalType", "arrow.duration-nanos")]),
5510 Field::new("duration_time_micros", DataType::Int64, false)
5511 .with_metadata([("logicalType", "arrow.duration-micros")]),
5512 Field::new("duration_time_millis", DataType::Int64, false)
5513 .with_metadata([("logicalType", "arrow.duration-millis")]),
5514 Field::new("duration_time_seconds", DataType::Int64, false)
5515 .with_metadata([("logicalType", "arrow.duration-seconds")]),
5516 ]));
5517
5518 let nanos =
5519 Arc::new(PrimitiveArray::<Int64Type>::from(vec![10, 20, 30, 40])) as ArrayRef;
5520 let micros = Arc::new(PrimitiveArray::<Int64Type>::from(vec![100, 200, 300, 400]))
5521 as ArrayRef;
5522 let millis = Arc::new(PrimitiveArray::<Int64Type>::from(vec![
5523 1000, 2000, 3000, 4000,
5524 ])) as ArrayRef;
5525 let seconds =
5526 Arc::new(PrimitiveArray::<Int64Type>::from(vec![1, 2, 3, 4])) as ArrayRef;
5527
5528 RecordBatch::try_new(schema, vec![nanos, micros, millis, seconds])?
5529 }
5530 };
5531
5532 assert_eq!(actual_batch, expected_batch);
5533
5534 Ok(())
5535 }
5536
5537 #[test]
5538 #[cfg(feature = "snappy")]
5540 fn test_dict_pages_offset_zero() {
5541 let file = arrow_test_data("avro/dict-page-offset-zero.avro");
5542 let batch = read_file(&file, 32, false);
5543 let num_rows = batch.num_rows();
5544 let expected_field = Int32Array::from(vec![Some(1552); num_rows]);
5545 let expected = RecordBatch::try_from_iter_with_nullable([(
5546 "l_partkey",
5547 Arc::new(expected_field) as Arc<dyn Array>,
5548 true,
5549 )])
5550 .unwrap();
5551 assert_eq!(batch, expected);
5552 }
5553
5554 #[test]
5555 #[cfg(feature = "snappy")]
5557 fn test_list_columns() {
5558 let file = arrow_test_data("avro/list_columns.avro");
5559 let mut int64_list_builder = ListBuilder::new(Int64Builder::new());
5560 {
5561 {
5562 let values = int64_list_builder.values();
5563 values.append_value(1);
5564 values.append_value(2);
5565 values.append_value(3);
5566 }
5567 int64_list_builder.append(true);
5568 }
5569 {
5570 {
5571 let values = int64_list_builder.values();
5572 values.append_null();
5573 values.append_value(1);
5574 }
5575 int64_list_builder.append(true);
5576 }
5577 {
5578 {
5579 let values = int64_list_builder.values();
5580 values.append_value(4);
5581 }
5582 int64_list_builder.append(true);
5583 }
5584 let int64_list = int64_list_builder.finish();
5585 let mut utf8_list_builder = ListBuilder::new(StringBuilder::new());
5586 {
5587 {
5588 let values = utf8_list_builder.values();
5589 values.append_value("abc");
5590 values.append_value("efg");
5591 values.append_value("hij");
5592 }
5593 utf8_list_builder.append(true);
5594 }
5595 {
5596 utf8_list_builder.append(false);
5597 }
5598 {
5599 {
5600 let values = utf8_list_builder.values();
5601 values.append_value("efg");
5602 values.append_null();
5603 values.append_value("hij");
5604 values.append_value("xyz");
5605 }
5606 utf8_list_builder.append(true);
5607 }
5608 let utf8_list = utf8_list_builder.finish();
5609 let expected = RecordBatch::try_from_iter_with_nullable([
5610 ("int64_list", Arc::new(int64_list) as Arc<dyn Array>, true),
5611 ("utf8_list", Arc::new(utf8_list) as Arc<dyn Array>, true),
5612 ])
5613 .unwrap();
5614 let batch = read_file(&file, 8, false);
5615 assert_eq!(batch, expected);
5616 }
5617
5618 #[test]
5619 #[cfg(feature = "snappy")]
5620 fn test_nested_lists() {
5621 use arrow_data::ArrayDataBuilder;
5622 let file = arrow_test_data("avro/nested_lists.snappy.avro");
5623 let inner_values = StringArray::from(vec![
5624 Some("a"),
5625 Some("b"),
5626 Some("c"),
5627 Some("d"),
5628 Some("a"),
5629 Some("b"),
5630 Some("c"),
5631 Some("d"),
5632 Some("e"),
5633 Some("a"),
5634 Some("b"),
5635 Some("c"),
5636 Some("d"),
5637 Some("e"),
5638 Some("f"),
5639 ]);
5640 let inner_offsets = Buffer::from_slice_ref([0, 2, 3, 3, 4, 6, 8, 8, 9, 11, 13, 14, 14, 15]);
5641 let inner_validity = [
5642 true, true, false, true, true, true, false, true, true, true, true, false, true,
5643 ];
5644 let inner_null_buffer = Buffer::from_iter(inner_validity.iter().copied());
5645 let inner_field = Field::new("item", DataType::Utf8, true);
5646 let inner_list_data = ArrayDataBuilder::new(DataType::List(Arc::new(inner_field)))
5647 .len(13)
5648 .add_buffer(inner_offsets)
5649 .add_child_data(inner_values.to_data())
5650 .null_bit_buffer(Some(inner_null_buffer))
5651 .build()
5652 .unwrap();
5653 let inner_list_array = ListArray::from(inner_list_data);
5654 let middle_offsets = Buffer::from_slice_ref([0, 2, 4, 6, 8, 11, 13]);
5655 let middle_validity = [true; 6];
5656 let middle_null_buffer = Buffer::from_iter(middle_validity.iter().copied());
5657 let middle_field = Field::new("item", inner_list_array.data_type().clone(), true);
5658 let middle_list_data = ArrayDataBuilder::new(DataType::List(Arc::new(middle_field)))
5659 .len(6)
5660 .add_buffer(middle_offsets)
5661 .add_child_data(inner_list_array.to_data())
5662 .null_bit_buffer(Some(middle_null_buffer))
5663 .build()
5664 .unwrap();
5665 let middle_list_array = ListArray::from(middle_list_data);
5666 let outer_offsets = Buffer::from_slice_ref([0, 2, 4, 6]);
5667 let outer_null_buffer = Buffer::from_slice_ref([0b111]); let outer_field = Field::new("item", middle_list_array.data_type().clone(), true);
5669 let outer_list_data = ArrayDataBuilder::new(DataType::List(Arc::new(outer_field)))
5670 .len(3)
5671 .add_buffer(outer_offsets)
5672 .add_child_data(middle_list_array.to_data())
5673 .null_bit_buffer(Some(outer_null_buffer))
5674 .build()
5675 .unwrap();
5676 let a_expected = ListArray::from(outer_list_data);
5677 let b_expected = Int32Array::from(vec![1, 1, 1]);
5678 let expected = RecordBatch::try_from_iter_with_nullable([
5679 ("a", Arc::new(a_expected) as Arc<dyn Array>, true),
5680 ("b", Arc::new(b_expected) as Arc<dyn Array>, true),
5681 ])
5682 .unwrap();
5683 let left = read_file(&file, 8, false);
5684 assert_eq!(left, expected, "Mismatch for batch size=8");
5685 let left_small = read_file(&file, 3, false);
5686 assert_eq!(left_small, expected, "Mismatch for batch size=3");
5687 }
5688
5689 #[test]
5690 fn test_simple() {
5691 let tests = [
5692 ("avro/simple_enum.avro", 4, build_expected_enum(), 2),
5693 ("avro/simple_fixed.avro", 2, build_expected_fixed(), 1),
5694 ];
5695
5696 fn build_expected_enum() -> RecordBatch {
5697 let keys_f1 = Int32Array::from(vec![0, 1, 2, 3]);
5699 let vals_f1 = StringArray::from(vec!["a", "b", "c", "d"]);
5700 let f1_dict =
5701 DictionaryArray::<Int32Type>::try_new(keys_f1, Arc::new(vals_f1)).unwrap();
5702 let keys_f2 = Int32Array::from(vec![2, 3, 0, 1]);
5703 let vals_f2 = StringArray::from(vec!["e", "f", "g", "h"]);
5704 let f2_dict =
5705 DictionaryArray::<Int32Type>::try_new(keys_f2, Arc::new(vals_f2)).unwrap();
5706 let keys_f3 = Int32Array::from(vec![Some(1), Some(2), None, Some(0)]);
5707 let vals_f3 = StringArray::from(vec!["i", "j", "k"]);
5708 let f3_dict =
5709 DictionaryArray::<Int32Type>::try_new(keys_f3, Arc::new(vals_f3)).unwrap();
5710 let dict_type =
5711 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8));
5712 let mut md_f1 = HashMap::new();
5713 md_f1.insert(
5714 AVRO_ENUM_SYMBOLS_METADATA_KEY.to_string(),
5715 r#"["a","b","c","d"]"#.to_string(),
5716 );
5717 md_f1.insert(AVRO_NAME_METADATA_KEY.to_string(), "enum1".to_string());
5718 md_f1.insert(AVRO_NAMESPACE_METADATA_KEY.to_string(), "ns1".to_string());
5719 let f1_field = Field::new("f1", dict_type.clone(), false).with_metadata(md_f1);
5720 let mut md_f2 = HashMap::new();
5721 md_f2.insert(
5722 AVRO_ENUM_SYMBOLS_METADATA_KEY.to_string(),
5723 r#"["e","f","g","h"]"#.to_string(),
5724 );
5725 md_f2.insert(AVRO_NAME_METADATA_KEY.to_string(), "enum2".to_string());
5726 md_f2.insert(AVRO_NAMESPACE_METADATA_KEY.to_string(), "ns2".to_string());
5727 let f2_field = Field::new("f2", dict_type.clone(), false).with_metadata(md_f2);
5728 let mut md_f3 = HashMap::new();
5729 md_f3.insert(
5730 AVRO_ENUM_SYMBOLS_METADATA_KEY.to_string(),
5731 r#"["i","j","k"]"#.to_string(),
5732 );
5733 md_f3.insert(AVRO_NAME_METADATA_KEY.to_string(), "enum3".to_string());
5734 md_f3.insert(AVRO_NAMESPACE_METADATA_KEY.to_string(), "ns1".to_string());
5735 let f3_field = Field::new("f3", dict_type.clone(), true).with_metadata(md_f3);
5736 let expected_schema = Arc::new(Schema::new(vec![f1_field, f2_field, f3_field]));
5737 RecordBatch::try_new(
5738 expected_schema,
5739 vec![
5740 Arc::new(f1_dict) as Arc<dyn Array>,
5741 Arc::new(f2_dict) as Arc<dyn Array>,
5742 Arc::new(f3_dict) as Arc<dyn Array>,
5743 ],
5744 )
5745 .unwrap()
5746 }
5747
5748 fn build_expected_fixed() -> RecordBatch {
5749 let f1 =
5750 FixedSizeBinaryArray::try_from_iter(vec![b"abcde", b"12345"].into_iter()).unwrap();
5751 let f2 =
5752 FixedSizeBinaryArray::try_from_iter(vec![b"fghijklmno", b"1234567890"].into_iter())
5753 .unwrap();
5754 let f3 = FixedSizeBinaryArray::try_from_sparse_iter_with_size(
5755 vec![Some(b"ABCDEF" as &[u8]), None].into_iter(),
5756 6,
5757 )
5758 .unwrap();
5759
5760 let mut md_f1 = HashMap::new();
5762 md_f1.insert(
5763 crate::schema::AVRO_NAME_METADATA_KEY.to_string(),
5764 "fixed1".to_string(),
5765 );
5766 md_f1.insert(
5767 crate::schema::AVRO_NAMESPACE_METADATA_KEY.to_string(),
5768 "ns1".to_string(),
5769 );
5770
5771 let mut md_f2 = HashMap::new();
5772 md_f2.insert(
5773 crate::schema::AVRO_NAME_METADATA_KEY.to_string(),
5774 "fixed2".to_string(),
5775 );
5776 md_f2.insert(
5777 crate::schema::AVRO_NAMESPACE_METADATA_KEY.to_string(),
5778 "ns2".to_string(),
5779 );
5780
5781 let mut md_f3 = HashMap::new();
5782 md_f3.insert(
5783 crate::schema::AVRO_NAME_METADATA_KEY.to_string(),
5784 "fixed3".to_string(),
5785 );
5786 md_f3.insert(
5787 crate::schema::AVRO_NAMESPACE_METADATA_KEY.to_string(),
5788 "ns1".to_string(),
5789 );
5790
5791 let expected_schema = Arc::new(Schema::new(vec![
5792 Field::new("f1", DataType::FixedSizeBinary(5), false).with_metadata(md_f1),
5793 Field::new("f2", DataType::FixedSizeBinary(10), false).with_metadata(md_f2),
5794 Field::new("f3", DataType::FixedSizeBinary(6), true).with_metadata(md_f3),
5795 ]));
5796
5797 RecordBatch::try_new(
5798 expected_schema,
5799 vec![
5800 Arc::new(f1) as Arc<dyn Array>,
5801 Arc::new(f2) as Arc<dyn Array>,
5802 Arc::new(f3) as Arc<dyn Array>,
5803 ],
5804 )
5805 .unwrap()
5806 }
5807 for (file_name, batch_size, expected, alt_batch_size) in tests {
5808 let file = arrow_test_data(file_name);
5809 let actual = read_file(&file, batch_size, false);
5810 assert_eq!(actual, expected);
5811 let actual2 = read_file(&file, alt_batch_size, false);
5812 assert_eq!(actual2, expected);
5813 }
5814 }
5815
5816 #[test]
5817 #[cfg(feature = "snappy")]
5818 fn test_single_nan() {
5819 let file = arrow_test_data("avro/single_nan.avro");
5820 let actual = read_file(&file, 1, false);
5821 use arrow_array::Float64Array;
5822 let schema = Arc::new(Schema::new(vec![Field::new(
5823 "mycol",
5824 DataType::Float64,
5825 true,
5826 )]));
5827 let col = Float64Array::from(vec![None]);
5828 let expected = RecordBatch::try_new(schema, vec![Arc::new(col)]).unwrap();
5829 assert_eq!(actual, expected);
5830 let actual2 = read_file(&file, 2, false);
5831 assert_eq!(actual2, expected);
5832 }
5833
5834 #[test]
5835 fn test_duration_uuid() {
5836 let batch = read_file("test/data/duration_uuid.avro", 4, false);
5837 let schema = batch.schema();
5838 let fields = schema.fields();
5839 assert_eq!(fields.len(), 2);
5840 assert_eq!(fields[0].name(), "duration_field");
5841 assert_eq!(
5842 fields[0].data_type(),
5843 &DataType::Interval(IntervalUnit::MonthDayNano)
5844 );
5845 assert_eq!(fields[1].name(), "uuid_field");
5846 assert_eq!(fields[1].data_type(), &DataType::FixedSizeBinary(16));
5847 assert_eq!(batch.num_rows(), 4);
5848 assert_eq!(batch.num_columns(), 2);
5849 let duration_array = batch
5850 .column(0)
5851 .as_any()
5852 .downcast_ref::<IntervalMonthDayNanoArray>()
5853 .unwrap();
5854 let expected_duration_array: IntervalMonthDayNanoArray = [
5855 Some(IntervalMonthDayNanoType::make_value(1, 15, 500_000_000)),
5856 Some(IntervalMonthDayNanoType::make_value(0, 5, 2_500_000_000)),
5857 Some(IntervalMonthDayNanoType::make_value(2, 0, 0)),
5858 Some(IntervalMonthDayNanoType::make_value(12, 31, 999_000_000)),
5859 ]
5860 .iter()
5861 .copied()
5862 .collect();
5863 assert_eq!(&expected_duration_array, duration_array);
5864 let uuid_array = batch
5865 .column(1)
5866 .as_any()
5867 .downcast_ref::<FixedSizeBinaryArray>()
5868 .unwrap();
5869 let expected_uuid_array = FixedSizeBinaryArray::try_from_sparse_iter_with_size(
5870 [
5871 Some([
5872 0xfe, 0x7b, 0xc3, 0x0b, 0x4c, 0xe8, 0x4c, 0x5e, 0xb6, 0x7c, 0x22, 0x34, 0xa2,
5873 0xd3, 0x8e, 0x66,
5874 ]),
5875 Some([
5876 0xb3, 0x3f, 0x2a, 0xd7, 0x97, 0xb4, 0x4d, 0xe1, 0x8b, 0xfe, 0x94, 0x94, 0x1d,
5877 0x60, 0x15, 0x6e,
5878 ]),
5879 Some([
5880 0x5f, 0x74, 0x92, 0x64, 0x07, 0x4b, 0x40, 0x05, 0x84, 0xbf, 0x11, 0x5e, 0xa8,
5881 0x4e, 0xd2, 0x0a,
5882 ]),
5883 Some([
5884 0x08, 0x26, 0xcc, 0x06, 0xd2, 0xe3, 0x45, 0x99, 0xb4, 0xad, 0xaf, 0x5f, 0xa6,
5885 0x90, 0x5c, 0xdb,
5886 ]),
5887 ]
5888 .into_iter(),
5889 16,
5890 )
5891 .unwrap();
5892 assert_eq!(&expected_uuid_array, uuid_array);
5893 }
5894
5895 #[test]
5896 #[cfg(feature = "snappy")]
5897 fn test_datapage_v2() {
5898 let file = arrow_test_data("avro/datapage_v2.snappy.avro");
5899 let batch = read_file(&file, 8, false);
5900 let a = StringArray::from(vec![
5901 Some("abc"),
5902 Some("abc"),
5903 Some("abc"),
5904 None,
5905 Some("abc"),
5906 ]);
5907 let b = Int32Array::from(vec![Some(1), Some(2), Some(3), Some(4), Some(5)]);
5908 let c = Float64Array::from(vec![Some(2.0), Some(3.0), Some(4.0), Some(5.0), Some(2.0)]);
5909 let d = BooleanArray::from(vec![
5910 Some(true),
5911 Some(true),
5912 Some(true),
5913 Some(false),
5914 Some(true),
5915 ]);
5916 let e_values = Int32Array::from(vec![
5917 Some(1),
5918 Some(2),
5919 Some(3),
5920 Some(1),
5921 Some(2),
5922 Some(3),
5923 Some(1),
5924 Some(2),
5925 ]);
5926 let e_offsets = OffsetBuffer::new(ScalarBuffer::from(vec![0i32, 3, 3, 3, 6, 8]));
5927 let e_validity = Some(NullBuffer::from(vec![true, false, false, true, true]));
5928 let field_e = Arc::new(Field::new("item", DataType::Int32, true));
5929 let e = ListArray::new(field_e, e_offsets, Arc::new(e_values), e_validity);
5930 let expected = RecordBatch::try_from_iter_with_nullable([
5931 ("a", Arc::new(a) as Arc<dyn Array>, true),
5932 ("b", Arc::new(b) as Arc<dyn Array>, true),
5933 ("c", Arc::new(c) as Arc<dyn Array>, true),
5934 ("d", Arc::new(d) as Arc<dyn Array>, true),
5935 ("e", Arc::new(e) as Arc<dyn Array>, true),
5936 ])
5937 .unwrap();
5938 assert_eq!(batch, expected);
5939 }
5940
5941 #[test]
5942 fn test_nested_records() {
5943 let f1_f1_1 = StringArray::from(vec!["aaa", "bbb"]);
5944 let f1_f1_2 = Int32Array::from(vec![10, 20]);
5945 let rounded_pi = (std::f64::consts::PI * 100.0).round() / 100.0;
5946 let f1_f1_3_1 = Float64Array::from(vec![rounded_pi, rounded_pi]);
5947 let f1_f1_3 = StructArray::from(vec![(
5948 Arc::new(Field::new("f1_3_1", DataType::Float64, false)),
5949 Arc::new(f1_f1_3_1) as Arc<dyn Array>,
5950 )]);
5951 let mut f1_3_md: HashMap<String, String> = HashMap::new();
5953 f1_3_md.insert(AVRO_NAMESPACE_METADATA_KEY.to_string(), "ns3".to_string());
5954 f1_3_md.insert(AVRO_NAME_METADATA_KEY.to_string(), "record3".to_string());
5955 let f1_expected = StructArray::from(vec![
5956 (
5957 Arc::new(Field::new("f1_1", DataType::Utf8, false)),
5958 Arc::new(f1_f1_1) as Arc<dyn Array>,
5959 ),
5960 (
5961 Arc::new(Field::new("f1_2", DataType::Int32, false)),
5962 Arc::new(f1_f1_2) as Arc<dyn Array>,
5963 ),
5964 (
5965 Arc::new(
5966 Field::new(
5967 "f1_3",
5968 DataType::Struct(Fields::from(vec![Field::new(
5969 "f1_3_1",
5970 DataType::Float64,
5971 false,
5972 )])),
5973 false,
5974 )
5975 .with_metadata(f1_3_md),
5976 ),
5977 Arc::new(f1_f1_3) as Arc<dyn Array>,
5978 ),
5979 ]);
5980 let f2_fields = [
5981 Field::new("f2_1", DataType::Boolean, false),
5982 Field::new("f2_2", DataType::Float32, false),
5983 ];
5984 let f2_struct_builder = StructBuilder::new(
5985 f2_fields
5986 .iter()
5987 .map(|f| Arc::new(f.clone()))
5988 .collect::<Vec<Arc<Field>>>(),
5989 vec![
5990 Box::new(BooleanBuilder::new()) as Box<dyn arrow_array::builder::ArrayBuilder>,
5991 Box::new(Float32Builder::new()) as Box<dyn arrow_array::builder::ArrayBuilder>,
5992 ],
5993 );
5994 let mut f2_list_builder = ListBuilder::new(f2_struct_builder);
5995 {
5996 let struct_builder = f2_list_builder.values();
5997 struct_builder.append(true);
5998 {
5999 let b = struct_builder.field_builder::<BooleanBuilder>(0).unwrap();
6000 b.append_value(true);
6001 }
6002 {
6003 let b = struct_builder.field_builder::<Float32Builder>(1).unwrap();
6004 b.append_value(1.2_f32);
6005 }
6006 struct_builder.append(true);
6007 {
6008 let b = struct_builder.field_builder::<BooleanBuilder>(0).unwrap();
6009 b.append_value(true);
6010 }
6011 {
6012 let b = struct_builder.field_builder::<Float32Builder>(1).unwrap();
6013 b.append_value(2.2_f32);
6014 }
6015 f2_list_builder.append(true);
6016 }
6017 {
6018 let struct_builder = f2_list_builder.values();
6019 struct_builder.append(true);
6020 {
6021 let b = struct_builder.field_builder::<BooleanBuilder>(0).unwrap();
6022 b.append_value(false);
6023 }
6024 {
6025 let b = struct_builder.field_builder::<Float32Builder>(1).unwrap();
6026 b.append_value(10.2_f32);
6027 }
6028 f2_list_builder.append(true);
6029 }
6030
6031 let list_array_with_nullable_items = f2_list_builder.finish();
6032 let mut f2_item_md: HashMap<String, String> = HashMap::new();
6034 f2_item_md.insert(AVRO_NAME_METADATA_KEY.to_string(), "record4".to_string());
6035 f2_item_md.insert(AVRO_NAMESPACE_METADATA_KEY.to_string(), "ns4".to_string());
6036 let item_field = Arc::new(
6037 Field::new(
6038 "item",
6039 list_array_with_nullable_items.values().data_type().clone(),
6040 false, )
6042 .with_metadata(f2_item_md),
6043 );
6044 let list_data_type = DataType::List(item_field);
6045 let f2_array_data = list_array_with_nullable_items
6046 .to_data()
6047 .into_builder()
6048 .data_type(list_data_type)
6049 .build()
6050 .unwrap();
6051 let f2_expected = ListArray::from(f2_array_data);
6052 let mut f3_struct_builder = StructBuilder::new(
6053 vec![Arc::new(Field::new("f3_1", DataType::Utf8, false))],
6054 vec![Box::new(StringBuilder::new()) as Box<dyn ArrayBuilder>],
6055 );
6056 f3_struct_builder.append(true);
6057 {
6058 let b = f3_struct_builder.field_builder::<StringBuilder>(0).unwrap();
6059 b.append_value("xyz");
6060 }
6061 f3_struct_builder.append(false);
6062 {
6063 let b = f3_struct_builder.field_builder::<StringBuilder>(0).unwrap();
6064 b.append_null();
6065 }
6066 let f3_expected = f3_struct_builder.finish();
6067 let f4_fields = [Field::new("f4_1", DataType::Int64, false)];
6068 let f4_struct_builder = StructBuilder::new(
6069 f4_fields
6070 .iter()
6071 .map(|f| Arc::new(f.clone()))
6072 .collect::<Vec<Arc<Field>>>(),
6073 vec![Box::new(Int64Builder::new()) as Box<dyn arrow_array::builder::ArrayBuilder>],
6074 );
6075 let mut f4_list_builder = ListBuilder::new(f4_struct_builder);
6076 {
6077 let struct_builder = f4_list_builder.values();
6078 struct_builder.append(true);
6079 {
6080 let b = struct_builder.field_builder::<Int64Builder>(0).unwrap();
6081 b.append_value(200);
6082 }
6083 struct_builder.append(false);
6084 {
6085 let b = struct_builder.field_builder::<Int64Builder>(0).unwrap();
6086 b.append_null();
6087 }
6088 f4_list_builder.append(true);
6089 }
6090 {
6091 let struct_builder = f4_list_builder.values();
6092 struct_builder.append(false);
6093 {
6094 let b = struct_builder.field_builder::<Int64Builder>(0).unwrap();
6095 b.append_null();
6096 }
6097 struct_builder.append(true);
6098 {
6099 let b = struct_builder.field_builder::<Int64Builder>(0).unwrap();
6100 b.append_value(300);
6101 }
6102 f4_list_builder.append(true);
6103 }
6104 let f4_expected = f4_list_builder.finish();
6105 let mut f4_item_md: HashMap<String, String> = HashMap::new();
6107 f4_item_md.insert(AVRO_NAMESPACE_METADATA_KEY.to_string(), "ns6".to_string());
6108 f4_item_md.insert(AVRO_NAME_METADATA_KEY.to_string(), "record6".to_string());
6109 let f4_item_field = Arc::new(
6110 Field::new("item", f4_expected.values().data_type().clone(), true)
6111 .with_metadata(f4_item_md),
6112 );
6113 let f4_list_data_type = DataType::List(f4_item_field);
6114 let f4_array_data = f4_expected
6115 .to_data()
6116 .into_builder()
6117 .data_type(f4_list_data_type)
6118 .build()
6119 .unwrap();
6120 let f4_expected = ListArray::from(f4_array_data);
6121 let mut f1_md: HashMap<String, String> = HashMap::new();
6123 f1_md.insert(AVRO_NAME_METADATA_KEY.to_string(), "record2".to_string());
6124 f1_md.insert(AVRO_NAMESPACE_METADATA_KEY.to_string(), "ns2".to_string());
6125 let mut f3_md: HashMap<String, String> = HashMap::new();
6126 f3_md.insert(AVRO_NAMESPACE_METADATA_KEY.to_string(), "ns5".to_string());
6127 f3_md.insert(AVRO_NAME_METADATA_KEY.to_string(), "record5".to_string());
6128 let expected_schema = Schema::new(vec![
6129 Field::new("f1", f1_expected.data_type().clone(), false).with_metadata(f1_md),
6130 Field::new("f2", f2_expected.data_type().clone(), false),
6131 Field::new("f3", f3_expected.data_type().clone(), true).with_metadata(f3_md),
6132 Field::new("f4", f4_expected.data_type().clone(), false),
6133 ]);
6134 let expected = RecordBatch::try_new(
6135 Arc::new(expected_schema),
6136 vec![
6137 Arc::new(f1_expected) as Arc<dyn Array>,
6138 Arc::new(f2_expected) as Arc<dyn Array>,
6139 Arc::new(f3_expected) as Arc<dyn Array>,
6140 Arc::new(f4_expected) as Arc<dyn Array>,
6141 ],
6142 )
6143 .unwrap();
6144 let file = arrow_test_data("avro/nested_records.avro");
6145 let batch_large = read_file(&file, 8, false);
6146 assert_eq!(
6147 batch_large, expected,
6148 "Decoded RecordBatch does not match expected data for nested records (batch size 8)"
6149 );
6150 let batch_small = read_file(&file, 3, false);
6151 assert_eq!(
6152 batch_small, expected,
6153 "Decoded RecordBatch does not match expected data for nested records (batch size 3)"
6154 );
6155 }
6156
6157 #[test]
6158 #[cfg(feature = "snappy")]
6160 fn test_repeated_no_annotation() {
6161 use arrow_data::ArrayDataBuilder;
6162 let file = arrow_test_data("avro/repeated_no_annotation.avro");
6163 let batch_large = read_file(&file, 8, false);
6164 let id_array = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
6166 let number_array = Int64Array::from(vec![
6168 Some(5555555555),
6169 Some(1111111111),
6170 Some(1111111111),
6171 Some(2222222222),
6172 Some(3333333333),
6173 ]);
6174 let kind_array =
6175 StringArray::from(vec![None, Some("home"), Some("home"), None, Some("mobile")]);
6176 let phone_fields = Fields::from(vec![
6177 Field::new("number", DataType::Int64, true),
6178 Field::new("kind", DataType::Utf8, true),
6179 ]);
6180 let phone_struct_data = ArrayDataBuilder::new(DataType::Struct(phone_fields))
6181 .len(5)
6182 .child_data(vec![number_array.into_data(), kind_array.into_data()])
6183 .build()
6184 .unwrap();
6185 let phone_struct_array = StructArray::from(phone_struct_data);
6186 let phone_list_offsets = Buffer::from_slice_ref([0i32, 0, 0, 0, 1, 2, 5]);
6188 let phone_list_validity = Buffer::from_iter([false, false, true, true, true, true]);
6189 let mut phone_item_md = HashMap::new();
6191 phone_item_md.insert(AVRO_NAME_METADATA_KEY.to_string(), "phone".to_string());
6192 phone_item_md.insert(
6193 AVRO_NAMESPACE_METADATA_KEY.to_string(),
6194 "topLevelRecord.phoneNumbers".to_string(),
6195 );
6196 let phone_item_field = Field::new("item", phone_struct_array.data_type().clone(), true)
6197 .with_metadata(phone_item_md);
6198 let phone_list_data = ArrayDataBuilder::new(DataType::List(Arc::new(phone_item_field)))
6199 .len(6)
6200 .add_buffer(phone_list_offsets)
6201 .null_bit_buffer(Some(phone_list_validity))
6202 .child_data(vec![phone_struct_array.into_data()])
6203 .build()
6204 .unwrap();
6205 let phone_list_array = ListArray::from(phone_list_data);
6206 let phone_numbers_validity = Buffer::from_iter([false, false, true, true, true, true]);
6208 let phone_numbers_field = Field::new("phone", phone_list_array.data_type().clone(), true);
6209 let phone_numbers_struct_data =
6210 ArrayDataBuilder::new(DataType::Struct(Fields::from(vec![phone_numbers_field])))
6211 .len(6)
6212 .null_bit_buffer(Some(phone_numbers_validity))
6213 .child_data(vec![phone_list_array.into_data()])
6214 .build()
6215 .unwrap();
6216 let phone_numbers_struct_array = StructArray::from(phone_numbers_struct_data);
6217 let mut phone_numbers_md = HashMap::new();
6219 phone_numbers_md.insert(
6220 AVRO_NAME_METADATA_KEY.to_string(),
6221 "phoneNumbers".to_string(),
6222 );
6223 phone_numbers_md.insert(
6224 AVRO_NAMESPACE_METADATA_KEY.to_string(),
6225 "topLevelRecord".to_string(),
6226 );
6227 let id_field = Field::new("id", DataType::Int32, true);
6228 let phone_numbers_schema_field = Field::new(
6229 "phoneNumbers",
6230 phone_numbers_struct_array.data_type().clone(),
6231 true,
6232 )
6233 .with_metadata(phone_numbers_md);
6234 let expected_schema = Schema::new(vec![id_field, phone_numbers_schema_field]);
6235 let expected = RecordBatch::try_new(
6237 Arc::new(expected_schema),
6238 vec![
6239 Arc::new(id_array) as _,
6240 Arc::new(phone_numbers_struct_array) as _,
6241 ],
6242 )
6243 .unwrap();
6244 assert_eq!(batch_large, expected, "Mismatch for batch_size=8");
6245 let batch_small = read_file(&file, 3, false);
6246 assert_eq!(batch_small, expected, "Mismatch for batch_size=3");
6247 }
6248
6249 #[test]
6250 #[cfg(feature = "snappy")]
6252 fn test_nonnullable_impala() {
6253 let file = arrow_test_data("avro/nonnullable.impala.avro");
6254 let id = Int64Array::from(vec![Some(8)]);
6255 let mut int_array_builder = ListBuilder::new(Int32Builder::new());
6256 {
6257 let vb = int_array_builder.values();
6258 vb.append_value(-1);
6259 }
6260 int_array_builder.append(true); let int_array = int_array_builder.finish();
6262 let mut iaa_builder = ListBuilder::new(ListBuilder::new(Int32Builder::new()));
6263 {
6264 let inner_list_builder = iaa_builder.values();
6265 {
6266 let vb = inner_list_builder.values();
6267 vb.append_value(-1);
6268 vb.append_value(-2);
6269 }
6270 inner_list_builder.append(true);
6271 inner_list_builder.append(true);
6272 }
6273 iaa_builder.append(true);
6274 let int_array_array = iaa_builder.finish();
6275 let field_names = MapFieldNames {
6276 entry: "entries".to_string(),
6277 key: "key".to_string(),
6278 value: "value".to_string(),
6279 };
6280 let mut int_map_builder =
6281 MapBuilder::new(Some(field_names), StringBuilder::new(), Int32Builder::new());
6282 {
6283 let (keys, vals) = int_map_builder.entries();
6284 keys.append_value("k1");
6285 vals.append_value(-1);
6286 }
6287 int_map_builder.append(true).unwrap(); let int_map = int_map_builder.finish();
6289 let field_names2 = MapFieldNames {
6290 entry: "entries".to_string(),
6291 key: "key".to_string(),
6292 value: "value".to_string(),
6293 };
6294 let mut ima_builder = ListBuilder::new(MapBuilder::new(
6295 Some(field_names2),
6296 StringBuilder::new(),
6297 Int32Builder::new(),
6298 ));
6299 {
6300 let map_builder = ima_builder.values();
6301 map_builder.append(true).unwrap();
6302 {
6303 let (keys, vals) = map_builder.entries();
6304 keys.append_value("k1");
6305 vals.append_value(1);
6306 }
6307 map_builder.append(true).unwrap();
6308 map_builder.append(true).unwrap();
6309 map_builder.append(true).unwrap();
6310 }
6311 ima_builder.append(true);
6312 let int_map_array_ = ima_builder.finish();
6313 let meta_nested_struct: HashMap<String, String> = [
6315 ("avro.name", "nested_Struct"),
6316 ("avro.namespace", "topLevelRecord"),
6317 ]
6318 .into_iter()
6319 .map(|(k, v)| (k.to_string(), v.to_string()))
6320 .collect();
6321 let meta_c: HashMap<String, String> = [
6322 ("avro.name", "c"),
6323 ("avro.namespace", "topLevelRecord.nested_Struct"),
6324 ]
6325 .into_iter()
6326 .map(|(k, v)| (k.to_string(), v.to_string()))
6327 .collect();
6328 let meta_d_item_struct: HashMap<String, String> = [
6329 ("avro.name", "D"),
6330 ("avro.namespace", "topLevelRecord.nested_Struct.c"),
6331 ]
6332 .into_iter()
6333 .map(|(k, v)| (k.to_string(), v.to_string()))
6334 .collect();
6335 let meta_g_value: HashMap<String, String> = [
6336 ("avro.name", "G"),
6337 ("avro.namespace", "topLevelRecord.nested_Struct"),
6338 ]
6339 .into_iter()
6340 .map(|(k, v)| (k.to_string(), v.to_string()))
6341 .collect();
6342 let meta_h: HashMap<String, String> = [
6343 ("avro.name", "h"),
6344 ("avro.namespace", "topLevelRecord.nested_Struct.G"),
6345 ]
6346 .into_iter()
6347 .map(|(k, v)| (k.to_string(), v.to_string()))
6348 .collect();
6349 let ef_struct_field = Arc::new(
6351 Field::new(
6352 "item",
6353 DataType::Struct(
6354 vec![
6355 Field::new("e", DataType::Int32, true),
6356 Field::new("f", DataType::Utf8, true),
6357 ]
6358 .into(),
6359 ),
6360 true,
6361 )
6362 .with_metadata(meta_d_item_struct.clone()),
6363 );
6364 let d_inner_list_field = Arc::new(Field::new(
6365 "item",
6366 DataType::List(ef_struct_field.clone()),
6367 true,
6368 ));
6369 let d_field = Field::new("D", DataType::List(d_inner_list_field.clone()), true);
6370 let i_list_field = Arc::new(Field::new("item", DataType::Float64, true));
6372 let i_field = Field::new("i", DataType::List(i_list_field.clone()), true);
6373 let h_field = Field::new("h", DataType::Struct(vec![i_field.clone()].into()), true)
6375 .with_metadata(meta_h.clone());
6376 let g_value_struct_field = Field::new(
6378 "value",
6379 DataType::Struct(vec![h_field.clone()].into()),
6380 true,
6381 )
6382 .with_metadata(meta_g_value.clone());
6383 let entries_struct_field = Field::new(
6385 "entries",
6386 DataType::Struct(
6387 vec![
6388 Field::new("key", DataType::Utf8, false),
6389 g_value_struct_field.clone(),
6390 ]
6391 .into(),
6392 ),
6393 false,
6394 );
6395 let a_field = Arc::new(Field::new("a", DataType::Int32, true));
6397 let b_field = Arc::new(Field::new(
6398 "B",
6399 DataType::List(Arc::new(Field::new("item", DataType::Int32, true))),
6400 true,
6401 ));
6402 let c_field = Arc::new(
6403 Field::new("c", DataType::Struct(vec![d_field.clone()].into()), true)
6404 .with_metadata(meta_c.clone()),
6405 );
6406 let g_field = Arc::new(Field::new(
6407 "G",
6408 DataType::Map(Arc::new(entries_struct_field.clone()), false),
6409 true,
6410 ));
6411 let mut nested_sb = StructBuilder::new(
6413 vec![
6414 a_field.clone(),
6415 b_field.clone(),
6416 c_field.clone(),
6417 g_field.clone(),
6418 ],
6419 vec![
6420 Box::new(Int32Builder::new()),
6421 Box::new(ListBuilder::new(Int32Builder::new())),
6422 {
6423 Box::new(StructBuilder::new(
6425 vec![Arc::new(d_field.clone())],
6426 vec![Box::new({
6427 let ef_struct_builder = StructBuilder::new(
6428 vec![
6429 Arc::new(Field::new("e", DataType::Int32, true)),
6430 Arc::new(Field::new("f", DataType::Utf8, true)),
6431 ],
6432 vec![
6433 Box::new(Int32Builder::new()),
6434 Box::new(StringBuilder::new()),
6435 ],
6436 );
6437 let list_of_ef = ListBuilder::new(ef_struct_builder)
6439 .with_field(ef_struct_field.clone());
6440 ListBuilder::new(list_of_ef)
6442 })],
6443 ))
6444 },
6445 {
6446 let map_field_names = MapFieldNames {
6447 entry: "entries".to_string(),
6448 key: "key".to_string(),
6449 value: "value".to_string(),
6450 };
6451 let i_list_builder = ListBuilder::new(Float64Builder::new());
6452 let h_struct_builder = StructBuilder::new(
6453 vec![Arc::new(Field::new(
6454 "i",
6455 DataType::List(i_list_field.clone()),
6456 true,
6457 ))],
6458 vec![Box::new(i_list_builder)],
6459 );
6460 let g_value_builder = StructBuilder::new(
6461 vec![Arc::new(
6462 Field::new("h", DataType::Struct(vec![i_field.clone()].into()), true)
6463 .with_metadata(meta_h.clone()),
6464 )],
6465 vec![Box::new(h_struct_builder)],
6466 );
6467 let map_builder = MapBuilder::new(
6469 Some(map_field_names),
6470 StringBuilder::new(),
6471 g_value_builder,
6472 )
6473 .with_values_field(Arc::new(
6474 Field::new(
6475 "value",
6476 DataType::Struct(vec![h_field.clone()].into()),
6477 true,
6478 )
6479 .with_metadata(meta_g_value.clone()),
6480 ));
6481
6482 Box::new(map_builder)
6483 },
6484 ],
6485 );
6486 nested_sb.append(true);
6487 {
6488 let a_builder = nested_sb.field_builder::<Int32Builder>(0).unwrap();
6489 a_builder.append_value(-1);
6490 }
6491 {
6492 let b_builder = nested_sb
6493 .field_builder::<ListBuilder<Int32Builder>>(1)
6494 .unwrap();
6495 {
6496 let vb = b_builder.values();
6497 vb.append_value(-1);
6498 }
6499 b_builder.append(true);
6500 }
6501 {
6502 let c_struct_builder = nested_sb.field_builder::<StructBuilder>(2).unwrap();
6503 c_struct_builder.append(true);
6504 let d_list_builder = c_struct_builder
6505 .field_builder::<ListBuilder<ListBuilder<StructBuilder>>>(0)
6506 .unwrap();
6507 {
6508 let sub_list_builder = d_list_builder.values();
6509 {
6510 let ef_struct = sub_list_builder.values();
6511 ef_struct.append(true);
6512 {
6513 let e_b = ef_struct.field_builder::<Int32Builder>(0).unwrap();
6514 e_b.append_value(-1);
6515 let f_b = ef_struct.field_builder::<StringBuilder>(1).unwrap();
6516 f_b.append_value("nonnullable");
6517 }
6518 sub_list_builder.append(true);
6519 }
6520 d_list_builder.append(true);
6521 }
6522 }
6523 {
6524 let g_map_builder = nested_sb
6525 .field_builder::<MapBuilder<StringBuilder, StructBuilder>>(3)
6526 .unwrap();
6527 g_map_builder.append(true).unwrap();
6528 }
6529 let nested_struct = nested_sb.finish();
6530 let schema = Arc::new(arrow_schema::Schema::new(vec![
6531 Field::new("ID", id.data_type().clone(), true),
6532 Field::new("Int_Array", int_array.data_type().clone(), true),
6533 Field::new("int_array_array", int_array_array.data_type().clone(), true),
6534 Field::new("Int_Map", int_map.data_type().clone(), true),
6535 Field::new("int_map_array", int_map_array_.data_type().clone(), true),
6536 Field::new("nested_Struct", nested_struct.data_type().clone(), true)
6537 .with_metadata(meta_nested_struct.clone()),
6538 ]));
6539 let expected = RecordBatch::try_new(
6540 schema,
6541 vec![
6542 Arc::new(id) as Arc<dyn Array>,
6543 Arc::new(int_array),
6544 Arc::new(int_array_array),
6545 Arc::new(int_map),
6546 Arc::new(int_map_array_),
6547 Arc::new(nested_struct),
6548 ],
6549 )
6550 .unwrap();
6551 let batch_large = read_file(&file, 8, false);
6552 assert_eq!(batch_large, expected, "Mismatch for batch_size=8");
6553 let batch_small = read_file(&file, 3, false);
6554 assert_eq!(batch_small, expected, "Mismatch for batch_size=3");
6555 }
6556
6557 #[test]
6558 fn test_nonnullable_impala_strict() {
6559 let file = arrow_test_data("avro/nonnullable.impala.avro");
6560 let err = read_file_strict(&file, 8, false).unwrap_err();
6561 assert!(err.to_string().contains(
6562 "Found Avro union of the form ['T','null'], which is disallowed in strict_mode"
6563 ));
6564 }
6565
6566 #[test]
6567 #[cfg(feature = "snappy")]
6569 fn test_nullable_impala() {
6570 let file = arrow_test_data("avro/nullable.impala.avro");
6571 let batch1 = read_file(&file, 3, false);
6572 let batch2 = read_file(&file, 8, false);
6573 assert_eq!(batch1, batch2);
6574 let batch = batch1;
6575 assert_eq!(batch.num_rows(), 7);
6576 let id_array = batch
6577 .column(0)
6578 .as_any()
6579 .downcast_ref::<Int64Array>()
6580 .expect("id column should be an Int64Array");
6581 let expected_ids = [1, 2, 3, 4, 5, 6, 7];
6582 for (i, &expected_id) in expected_ids.iter().enumerate() {
6583 assert_eq!(id_array.value(i), expected_id, "Mismatch in id at row {i}",);
6584 }
6585 let int_array = batch
6586 .column(1)
6587 .as_any()
6588 .downcast_ref::<ListArray>()
6589 .expect("int_array column should be a ListArray");
6590 {
6591 let offsets = int_array.value_offsets();
6592 let start = offsets[0] as usize;
6593 let end = offsets[1] as usize;
6594 let values = int_array
6595 .values()
6596 .as_any()
6597 .downcast_ref::<Int32Array>()
6598 .expect("Values of int_array should be an Int32Array");
6599 let row0: Vec<Option<i32>> = (start..end).map(|i| Some(values.value(i))).collect();
6600 assert_eq!(
6601 row0,
6602 vec![Some(1), Some(2), Some(3)],
6603 "Mismatch in int_array row 0"
6604 );
6605 }
6606 let nested_struct = batch
6607 .column(5)
6608 .as_any()
6609 .downcast_ref::<StructArray>()
6610 .expect("nested_struct column should be a StructArray");
6611 let a_array = nested_struct
6612 .column_by_name("A")
6613 .expect("Field A should exist in nested_struct")
6614 .as_any()
6615 .downcast_ref::<Int32Array>()
6616 .expect("Field A should be an Int32Array");
6617 assert_eq!(a_array.value(0), 1, "Mismatch in nested_struct.A at row 0");
6618 assert!(
6619 !a_array.is_valid(1),
6620 "Expected null in nested_struct.A at row 1"
6621 );
6622 assert!(
6623 !a_array.is_valid(3),
6624 "Expected null in nested_struct.A at row 3"
6625 );
6626 assert_eq!(a_array.value(6), 7, "Mismatch in nested_struct.A at row 6");
6627 }
6628
6629 #[test]
6630 fn test_nullable_impala_strict() {
6631 let file = arrow_test_data("avro/nullable.impala.avro");
6632 let err = read_file_strict(&file, 8, false).unwrap_err();
6633 assert!(err.to_string().contains(
6634 "Found Avro union of the form ['T','null'], which is disallowed in strict_mode"
6635 ));
6636 }
6637
6638 #[test]
6639 fn test_nested_record_type_reuse() {
6640 let batch = read_file("test/data/nested_record_reuse.avro", 8, false);
6666 let schema = batch.schema();
6667
6668 assert_eq!(schema.fields().len(), 3);
6670 let fields = schema.fields();
6671 assert_eq!(fields[0].name(), "nested");
6672 assert_eq!(fields[1].name(), "nestedRecord");
6673 assert_eq!(fields[2].name(), "nestedArray");
6674 assert!(matches!(fields[0].data_type(), DataType::Struct(_)));
6675 assert!(matches!(fields[1].data_type(), DataType::Struct(_)));
6676 assert!(matches!(fields[2].data_type(), DataType::List(_)));
6677
6678 if let DataType::Struct(nested_fields) = fields[0].data_type() {
6680 assert_eq!(nested_fields.len(), 1);
6681 assert_eq!(nested_fields[0].name(), "nested_int");
6682 assert_eq!(nested_fields[0].data_type(), &DataType::Int32);
6683 }
6684
6685 assert_eq!(fields[0].data_type(), fields[1].data_type());
6687 if let DataType::List(array_field) = fields[2].data_type() {
6688 assert_eq!(array_field.data_type(), fields[0].data_type());
6689 }
6690
6691 assert_eq!(batch.num_rows(), 2);
6693 assert_eq!(batch.num_columns(), 3);
6694
6695 let nested_col = batch
6697 .column(0)
6698 .as_any()
6699 .downcast_ref::<StructArray>()
6700 .unwrap();
6701 let nested_int_array = nested_col
6702 .column_by_name("nested_int")
6703 .unwrap()
6704 .as_any()
6705 .downcast_ref::<Int32Array>()
6706 .unwrap();
6707 assert_eq!(nested_int_array.value(0), 42);
6708 assert_eq!(nested_int_array.value(1), 99);
6709
6710 let nested_record_col = batch
6712 .column(1)
6713 .as_any()
6714 .downcast_ref::<StructArray>()
6715 .unwrap();
6716 let nested_record_int_array = nested_record_col
6717 .column_by_name("nested_int")
6718 .unwrap()
6719 .as_any()
6720 .downcast_ref::<Int32Array>()
6721 .unwrap();
6722 assert_eq!(nested_record_int_array.value(0), 100);
6723 assert_eq!(nested_record_int_array.value(1), 200);
6724
6725 let nested_array_col = batch
6727 .column(2)
6728 .as_any()
6729 .downcast_ref::<ListArray>()
6730 .unwrap();
6731 assert_eq!(nested_array_col.len(), 2);
6732 let first_array_struct = nested_array_col.value(0);
6733 let first_array_struct_array = first_array_struct
6734 .as_any()
6735 .downcast_ref::<StructArray>()
6736 .unwrap();
6737 let first_array_int_values = first_array_struct_array
6738 .column_by_name("nested_int")
6739 .unwrap()
6740 .as_any()
6741 .downcast_ref::<Int32Array>()
6742 .unwrap();
6743 assert_eq!(first_array_int_values.len(), 3);
6744 assert_eq!(first_array_int_values.value(0), 1);
6745 assert_eq!(first_array_int_values.value(1), 2);
6746 assert_eq!(first_array_int_values.value(2), 3);
6747 }
6748
6749 #[test]
6750 fn test_enum_type_reuse() {
6751 let batch = read_file("test/data/enum_reuse.avro", 8, false);
6774 let schema = batch.schema();
6775
6776 assert_eq!(schema.fields().len(), 3);
6778 let fields = schema.fields();
6779 assert_eq!(fields[0].name(), "status");
6780 assert_eq!(fields[1].name(), "backupStatus");
6781 assert_eq!(fields[2].name(), "statusHistory");
6782 assert!(matches!(fields[0].data_type(), DataType::Dictionary(_, _)));
6783 assert!(matches!(fields[1].data_type(), DataType::Dictionary(_, _)));
6784 assert!(matches!(fields[2].data_type(), DataType::List(_)));
6785
6786 if let DataType::Dictionary(key_type, value_type) = fields[0].data_type() {
6787 assert_eq!(key_type.as_ref(), &DataType::Int32);
6788 assert_eq!(value_type.as_ref(), &DataType::Utf8);
6789 }
6790
6791 assert_eq!(fields[0].data_type(), fields[1].data_type());
6793 if let DataType::List(array_field) = fields[2].data_type() {
6794 assert_eq!(array_field.data_type(), fields[0].data_type());
6795 }
6796
6797 assert_eq!(batch.num_rows(), 2);
6799 assert_eq!(batch.num_columns(), 3);
6800
6801 let status_col = batch
6803 .column(0)
6804 .as_any()
6805 .downcast_ref::<DictionaryArray<Int32Type>>()
6806 .unwrap();
6807 let status_values = status_col
6808 .values()
6809 .as_any()
6810 .downcast_ref::<StringArray>()
6811 .unwrap();
6812
6813 assert_eq!(status_values.value(status_col.key(0).unwrap()), "ACTIVE");
6815 assert_eq!(status_values.value(status_col.key(1).unwrap()), "PENDING");
6816
6817 let backup_status_col = batch
6819 .column(1)
6820 .as_any()
6821 .downcast_ref::<DictionaryArray<Int32Type>>()
6822 .unwrap();
6823 let backup_status_values = backup_status_col
6824 .values()
6825 .as_any()
6826 .downcast_ref::<StringArray>()
6827 .unwrap();
6828
6829 assert_eq!(
6831 backup_status_values.value(backup_status_col.key(0).unwrap()),
6832 "INACTIVE"
6833 );
6834 assert_eq!(
6835 backup_status_values.value(backup_status_col.key(1).unwrap()),
6836 "ACTIVE"
6837 );
6838
6839 let status_history_col = batch
6841 .column(2)
6842 .as_any()
6843 .downcast_ref::<ListArray>()
6844 .unwrap();
6845 assert_eq!(status_history_col.len(), 2);
6846
6847 let first_array_dict = status_history_col.value(0);
6849 let first_array_dict_array = first_array_dict
6850 .as_any()
6851 .downcast_ref::<DictionaryArray<Int32Type>>()
6852 .unwrap();
6853 let first_array_values = first_array_dict_array
6854 .values()
6855 .as_any()
6856 .downcast_ref::<StringArray>()
6857 .unwrap();
6858
6859 assert_eq!(first_array_dict_array.len(), 3);
6861 assert_eq!(
6862 first_array_values.value(first_array_dict_array.key(0).unwrap()),
6863 "PENDING"
6864 );
6865 assert_eq!(
6866 first_array_values.value(first_array_dict_array.key(1).unwrap()),
6867 "ACTIVE"
6868 );
6869 assert_eq!(
6870 first_array_values.value(first_array_dict_array.key(2).unwrap()),
6871 "INACTIVE"
6872 );
6873 }
6874
6875 #[test]
6876 fn test_bad_varint_bug_nullable_array_items() {
6877 use flate2::read::GzDecoder;
6878 use std::io::Read;
6879 let manifest_dir = env!("CARGO_MANIFEST_DIR");
6880 let gz_path = format!("{manifest_dir}/test/data/bad-varint-bug.avro.gz");
6881 let gz_file = File::open(&gz_path).expect("test file should exist");
6882 let mut decoder = GzDecoder::new(gz_file);
6883 let mut avro_bytes = Vec::new();
6884 decoder
6885 .read_to_end(&mut avro_bytes)
6886 .expect("should decompress");
6887 let reader_arrow_schema = Schema::new(vec![Field::new(
6888 "int_array",
6889 DataType::List(Arc::new(Field::new("element", DataType::Int32, true))),
6890 true,
6891 )])
6892 .with_metadata(HashMap::from([("avro.name".into(), "table".into())]));
6893 let reader_schema = AvroSchema::try_from(&reader_arrow_schema)
6894 .expect("should convert Arrow schema to Avro");
6895 let mut reader = ReaderBuilder::new()
6896 .with_reader_schema(reader_schema)
6897 .build(Cursor::new(avro_bytes))
6898 .expect("should build reader");
6899 let batch = reader
6900 .next()
6901 .expect("should have one batch")
6902 .expect("reading should succeed without bad varint error");
6903 assert_eq!(batch.num_rows(), 1);
6904 let list_col = batch
6905 .column(0)
6906 .as_any()
6907 .downcast_ref::<ListArray>()
6908 .expect("should be ListArray");
6909 assert_eq!(list_col.len(), 1);
6910 let values = list_col.values();
6911 let int_values = values.as_primitive::<Int32Type>();
6912 assert_eq!(int_values.len(), 2);
6913 assert_eq!(int_values.value(0), 1);
6914 assert_eq!(int_values.value(1), 2);
6915 }
6916
6917 #[test]
6918 fn test_nested_record_field_addition() {
6919 let file = arrow_test_data("avro/nested_records.avro");
6920
6921 let reader_schema = AvroSchema::new(
6933 r#"
6934 {
6935 "type": "record",
6936 "name": "record1",
6937 "namespace": "ns1",
6938 "fields": [
6939 {
6940 "name": "f1",
6941 "type": [
6942 "null",
6943 {
6944 "type": "record",
6945 "name": "record2",
6946 "namespace": "ns2",
6947 "fields": [
6948 {
6949 "name": "f1_1",
6950 "type": "string"
6951 },
6952 {
6953 "name": "f1_2",
6954 "type": "int"
6955 },
6956 {
6957 "name": "f1_3",
6958 "type": {
6959 "type": "record",
6960 "name": "record3",
6961 "namespace": "ns3",
6962 "fields": [
6963 {
6964 "name": "f1_3_1",
6965 "type": "double"
6966 }
6967 ]
6968 }
6969 },
6970 {
6971 "name": "f1_4",
6972 "type": ["null", "int"],
6973 "default": null
6974 }
6975 ]
6976 }
6977 ]
6978 },
6979 {
6980 "name": "f2",
6981 "type": {
6982 "type": "array",
6983 "items": {
6984 "type": "record",
6985 "name": "record4",
6986 "namespace": "ns4",
6987 "fields": [
6988 {
6989 "name": "f2_1",
6990 "type": "boolean"
6991 },
6992 {
6993 "name": "f2_2",
6994 "type": "float"
6995 },
6996 {
6997 "name": "f2_3",
6998 "type": ["null", "int"],
6999 "default": 42
7000 }
7001 ]
7002 }
7003 }
7004 },
7005 {
7006 "name": "f3",
7007 "type": [
7008 "null",
7009 {
7010 "type": "record",
7011 "name": "record5",
7012 "namespace": "ns5",
7013 "fields": [
7014 {
7015 "name": "f3_0",
7016 "type": "string",
7017 "default": "lorem ipsum"
7018 },
7019 {
7020 "name": "f3_1",
7021 "type": "string"
7022 }
7023 ]
7024 }
7025 ],
7026 "default": null
7027 },
7028 {
7029 "name": "f4",
7030 "type": {
7031 "type": "array",
7032 "items": [
7033 "null",
7034 {
7035 "type": "record",
7036 "name": "record6",
7037 "namespace": "ns6",
7038 "fields": [
7039 {
7040 "name": "f4_1",
7041 "type": "long"
7042 }
7043 ]
7044 }
7045 ]
7046 }
7047 }
7048 ]
7049 }
7050 "#
7051 .to_string(),
7052 );
7053
7054 let file = File::open(&file).unwrap();
7055 let mut reader = ReaderBuilder::new()
7056 .with_reader_schema(reader_schema)
7057 .build(BufReader::new(file))
7058 .expect("reader with evolved reader schema should be built successfully");
7059
7060 let batch = reader
7061 .next()
7062 .expect("should have at least one batch")
7063 .expect("reading should succeed");
7064
7065 assert!(batch.num_rows() > 0);
7066
7067 let schema = batch.schema();
7068
7069 let f1_field = schema.field_with_name("f1").expect("f1 field should exist");
7070 if let DataType::Struct(f1_fields) = f1_field.data_type() {
7071 let (_, f1_4) = f1_fields
7072 .find("f1_4")
7073 .expect("f1_4 field should be present in record2");
7074 assert!(f1_4.is_nullable(), "f1_4 should be nullable");
7075 assert_eq!(f1_4.data_type(), &DataType::Int32, "f1_4 should be Int32");
7076 assert_eq!(
7077 f1_4.metadata().get("avro.field.default"),
7078 Some(&"null".to_string()),
7079 "f1_4 should have null default value in metadata"
7080 );
7081 } else {
7082 panic!("f1 should be a struct");
7083 }
7084
7085 let f2_field = schema.field_with_name("f2").expect("f2 field should exist");
7086 if let DataType::List(f2_items_field) = f2_field.data_type() {
7087 if let DataType::Struct(f2_items_fields) = f2_items_field.data_type() {
7088 let (_, f2_3) = f2_items_fields
7089 .find("f2_3")
7090 .expect("f2_3 field should be present in record4");
7091 assert!(f2_3.is_nullable(), "f2_3 should be nullable");
7092 assert_eq!(f2_3.data_type(), &DataType::Int32, "f2_3 should be Int32");
7093 assert_eq!(
7094 f2_3.metadata().get("avro.field.default"),
7095 Some(&"42".to_string()),
7096 "f2_3 should have 42 default value in metadata"
7097 );
7098 } else {
7099 panic!("f2 array items should be a struct");
7100 }
7101 } else {
7102 panic!("f2 should be a list");
7103 }
7104
7105 let f3_field = schema.field_with_name("f3").expect("f3 field should exist");
7106 assert!(f3_field.is_nullable(), "f3 should be nullable");
7107 if let DataType::Struct(f3_fields) = f3_field.data_type() {
7108 let (_, f3_0) = f3_fields
7109 .find("f3_0")
7110 .expect("f3_0 field should be present in record5");
7111 assert!(!f3_0.is_nullable(), "f3_0 should be non-nullable");
7112 assert_eq!(f3_0.data_type(), &DataType::Utf8, "f3_0 should be a string");
7113 assert_eq!(
7114 f3_0.metadata().get("avro.field.default"),
7115 Some(&"\"lorem ipsum\"".to_string()),
7116 "f3_0 should have \"lorem ipsum\" default value in metadata"
7117 );
7118 } else {
7119 panic!("f3 should be a struct");
7120 }
7121
7122 let num_rows = batch.num_rows();
7124
7125 let f1_array = batch
7127 .column_by_name("f1")
7128 .expect("f1 column should exist")
7129 .as_struct();
7130 let f1_4_array = f1_array
7131 .column_by_name("f1_4")
7132 .expect("f1_4 column should exist in f1 struct")
7133 .as_primitive::<Int32Type>();
7134
7135 assert_eq!(f1_4_array.null_count(), num_rows);
7136
7137 let f2_array = batch
7138 .column_by_name("f2")
7139 .expect("f2 column should exist")
7140 .as_list::<i32>();
7141
7142 for i in 0..num_rows {
7143 assert!(!f2_array.is_null(i));
7144 let f2_value = f2_array.value(i);
7145 let f2_record_array = f2_value.as_struct();
7146 let f2_3_array = f2_record_array
7147 .column_by_name("f2_3")
7148 .expect("f2_3 column should exist in f2 array items")
7149 .as_primitive::<Int32Type>();
7150
7151 for j in 0..f2_3_array.len() {
7152 assert!(!f2_3_array.is_null(j));
7153 assert_eq!(f2_3_array.value(j), 42);
7154 }
7155 }
7156
7157 let f3_array = batch
7158 .column_by_name("f3")
7159 .expect("f3 column should exist")
7160 .as_struct();
7161 let f3_0_array = f3_array
7162 .column_by_name("f3_0")
7163 .expect("f3_0 column should exist in f3 struct")
7164 .as_string::<i32>();
7165
7166 for i in 0..num_rows {
7167 if !f3_array.is_null(i) {
7169 assert!(!f3_0_array.is_null(i));
7170 assert_eq!(f3_0_array.value(i), "lorem ipsum");
7171 }
7172 }
7173 }
7174
7175 fn corrupt_first_block_payload_byte(
7176 mut bytes: Vec<u8>,
7177 field_offset: usize,
7178 expected_original: u8,
7179 replacement: u8,
7180 ) -> Vec<u8> {
7181 let mut header_decoder = HeaderDecoder::default();
7182 let header_len = header_decoder.decode(&bytes).expect("decode header");
7183 assert!(header_decoder.flush().is_some(), "decode complete header");
7184
7185 let mut cursor = &bytes[header_len..];
7186 let (_, count_len) = crate::reader::vlq::read_varint(cursor).expect("decode block count");
7187 cursor = &cursor[count_len..];
7188 let (_, size_len) = crate::reader::vlq::read_varint(cursor).expect("decode block size");
7189 let data_start = header_len + count_len + size_len;
7190 let target = data_start + field_offset;
7191
7192 assert!(
7193 target < bytes.len(),
7194 "target byte offset {target} out of bounds for input length {}",
7195 bytes.len()
7196 );
7197 assert_eq!(
7198 bytes[target], expected_original,
7199 "unexpected original byte at payload offset {field_offset}"
7200 );
7201 bytes[target] = replacement;
7202 bytes
7203 }
7204
7205 #[test]
7206 fn ocf_projection_rejects_overflowing_varint_in_skipped_long_field() {
7207 let writer_schema = Schema::new(vec![
7211 Field::new("bad_long", DataType::Int64, false),
7212 Field::new("keep", DataType::Int32, false),
7213 ]);
7214 let batch = RecordBatch::try_new(
7215 Arc::new(writer_schema.clone()),
7216 vec![
7217 Arc::new(Int64Array::from(vec![i64::MIN])) as ArrayRef,
7218 Arc::new(Int32Array::from(vec![7])) as ArrayRef,
7219 ],
7220 )
7221 .expect("build writer batch");
7222 let bytes = write_ocf(&writer_schema, &[batch]);
7223 let mutated = corrupt_first_block_payload_byte(bytes, 9, 0x01, 0x02);
7224
7225 let err = ReaderBuilder::new()
7226 .build(Cursor::new(mutated.clone()))
7227 .expect("build full reader")
7228 .collect::<Result<Vec<_>, _>>()
7229 .expect_err("full decode should reject malformed varint");
7230 assert!(matches!(err, ArrowError::AvroError(_)));
7231 assert!(err.to_string().contains("bad varint"));
7232
7233 let err = ReaderBuilder::new()
7234 .with_projection(vec![1])
7235 .build(Cursor::new(mutated))
7236 .expect("build projected reader")
7237 .collect::<Result<Vec<_>, _>>()
7238 .expect_err("projection must also reject malformed skipped varint");
7239 assert!(matches!(err, ArrowError::AvroError(_)));
7240 assert!(err.to_string().contains("bad varint"));
7241 }
7242
7243 #[test]
7244 fn ocf_projection_rejects_i32_overflow_in_skipped_int_field() {
7245 let writer_schema = Schema::new(vec![
7249 Field::new("bad_int", DataType::Int32, false),
7250 Field::new("keep", DataType::Int64, false),
7251 ]);
7252 let batch = RecordBatch::try_new(
7253 Arc::new(writer_schema.clone()),
7254 vec![
7255 Arc::new(Int32Array::from(vec![i32::MIN])) as ArrayRef,
7256 Arc::new(Int64Array::from(vec![11])) as ArrayRef,
7257 ],
7258 )
7259 .expect("build writer batch");
7260 let bytes = write_ocf(&writer_schema, &[batch]);
7261 let mutated = corrupt_first_block_payload_byte(bytes, 4, 0x0f, 0x10);
7262
7263 let err = ReaderBuilder::new()
7264 .build(Cursor::new(mutated.clone()))
7265 .expect("build full reader")
7266 .collect::<Result<Vec<_>, _>>()
7267 .expect_err("full decode should reject int overflow");
7268 assert!(matches!(err, ArrowError::AvroError(_)));
7269 assert!(err.to_string().contains("varint overflow"));
7270
7271 let err = ReaderBuilder::new()
7272 .with_projection(vec![1])
7273 .build(Cursor::new(mutated))
7274 .expect("build projected reader")
7275 .collect::<Result<Vec<_>, _>>()
7276 .expect_err("projection must also reject skipped int overflow");
7277 assert!(matches!(err, ArrowError::AvroError(_)));
7278 assert!(err.to_string().contains("varint overflow"));
7279 }
7280
7281 #[test]
7282 fn comprehensive_e2e_test() {
7283 let path = "test/data/comprehensive_e2e.avro";
7284 let batch = read_file(path, 1024, false);
7285 let schema = batch.schema();
7286
7287 #[inline]
7288 fn tid_by_name(fields: &UnionFields, want: &str) -> i8 {
7289 for (tid, f) in fields.iter() {
7290 if f.name() == want {
7291 return tid;
7292 }
7293 }
7294 panic!("union child '{want}' not found");
7295 }
7296
7297 #[inline]
7298 fn tid_by_dt(fields: &UnionFields, pred: impl Fn(&DataType) -> bool) -> i8 {
7299 for (tid, f) in fields.iter() {
7300 if pred(f.data_type()) {
7301 return tid;
7302 }
7303 }
7304 panic!("no union child matches predicate");
7305 }
7306
7307 fn mk_dense_union(
7308 fields: &UnionFields,
7309 type_ids: Vec<i8>,
7310 offsets: Vec<i32>,
7311 provide: impl Fn(&Field) -> Option<ArrayRef>,
7312 ) -> ArrayRef {
7313 fn empty_child_for(dt: &DataType) -> Arc<dyn Array> {
7314 match dt {
7315 DataType::Null => Arc::new(NullArray::new(0)),
7316 DataType::Boolean => Arc::new(BooleanArray::from(Vec::<bool>::new())),
7317 DataType::Int32 => Arc::new(Int32Array::from(Vec::<i32>::new())),
7318 DataType::Int64 => Arc::new(Int64Array::from(Vec::<i64>::new())),
7319 DataType::Float32 => Arc::new(Float32Array::from(Vec::<f32>::new())),
7320 DataType::Float64 => Arc::new(Float64Array::from(Vec::<f64>::new())),
7321 DataType::Binary => Arc::new(BinaryArray::from(Vec::<&[u8]>::new())),
7322 DataType::Utf8 => Arc::new(StringArray::from(Vec::<&str>::new())),
7323 DataType::Date32 => Arc::new(Date32Array::from(Vec::<i32>::new())),
7324 DataType::Time32(arrow_schema::TimeUnit::Millisecond) => {
7325 Arc::new(Time32MillisecondArray::from(Vec::<i32>::new()))
7326 }
7327 DataType::Time64(arrow_schema::TimeUnit::Microsecond) => {
7328 Arc::new(Time64MicrosecondArray::from(Vec::<i64>::new()))
7329 }
7330 DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, tz) => {
7331 let a = TimestampMillisecondArray::from(Vec::<i64>::new());
7332 Arc::new(if let Some(tz) = tz {
7333 a.with_timezone(tz.clone())
7334 } else {
7335 a
7336 })
7337 }
7338 DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, tz) => {
7339 let a = TimestampMicrosecondArray::from(Vec::<i64>::new());
7340 Arc::new(if let Some(tz) = tz {
7341 a.with_timezone(tz.clone())
7342 } else {
7343 a
7344 })
7345 }
7346 DataType::Interval(IntervalUnit::MonthDayNano) => Arc::new(
7347 IntervalMonthDayNanoArray::from(Vec::<IntervalMonthDayNano>::new()),
7348 ),
7349 DataType::FixedSizeBinary(sz) => Arc::new(
7350 FixedSizeBinaryArray::try_from_sparse_iter_with_size(
7351 std::iter::empty::<Option<Vec<u8>>>(),
7352 *sz,
7353 )
7354 .unwrap(),
7355 ),
7356 DataType::Dictionary(_, _) => {
7357 let keys = Int32Array::from(Vec::<i32>::new());
7358 let values = Arc::new(StringArray::from(Vec::<&str>::new()));
7359 Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap())
7360 }
7361 DataType::Struct(fields) => {
7362 let children: Vec<ArrayRef> = fields
7363 .iter()
7364 .map(|f| empty_child_for(f.data_type()) as ArrayRef)
7365 .collect();
7366 Arc::new(StructArray::new(fields.clone(), children, None))
7367 }
7368 DataType::List(field) => {
7369 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0]));
7370 Arc::new(
7371 ListArray::try_new(
7372 field.clone(),
7373 offsets,
7374 empty_child_for(field.data_type()),
7375 None,
7376 )
7377 .unwrap(),
7378 )
7379 }
7380 DataType::Map(entry_field, is_sorted) => {
7381 let (key_field, val_field) = match entry_field.data_type() {
7382 DataType::Struct(fs) => (fs[0].clone(), fs[1].clone()),
7383 other => panic!("unexpected map entries type: {other:?}"),
7384 };
7385 let keys = StringArray::from(Vec::<&str>::new());
7386 let vals: ArrayRef = match val_field.data_type() {
7387 DataType::Null => Arc::new(NullArray::new(0)) as ArrayRef,
7388 DataType::Boolean => {
7389 Arc::new(BooleanArray::from(Vec::<bool>::new())) as ArrayRef
7390 }
7391 DataType::Int32 => {
7392 Arc::new(Int32Array::from(Vec::<i32>::new())) as ArrayRef
7393 }
7394 DataType::Int64 => {
7395 Arc::new(Int64Array::from(Vec::<i64>::new())) as ArrayRef
7396 }
7397 DataType::Float32 => {
7398 Arc::new(Float32Array::from(Vec::<f32>::new())) as ArrayRef
7399 }
7400 DataType::Float64 => {
7401 Arc::new(Float64Array::from(Vec::<f64>::new())) as ArrayRef
7402 }
7403 DataType::Utf8 => {
7404 Arc::new(StringArray::from(Vec::<&str>::new())) as ArrayRef
7405 }
7406 DataType::Binary => {
7407 Arc::new(BinaryArray::from(Vec::<&[u8]>::new())) as ArrayRef
7408 }
7409 DataType::Union(uf, _) => {
7410 let children: Vec<ArrayRef> = uf
7411 .iter()
7412 .map(|(_, f)| empty_child_for(f.data_type()))
7413 .collect();
7414 Arc::new(
7415 UnionArray::try_new(
7416 uf.clone(),
7417 ScalarBuffer::<i8>::from(Vec::<i8>::new()),
7418 Some(ScalarBuffer::<i32>::from(Vec::<i32>::new())),
7419 children,
7420 )
7421 .unwrap(),
7422 ) as ArrayRef
7423 }
7424 other => panic!("unsupported map value type: {other:?}"),
7425 };
7426 let entries = StructArray::new(
7427 Fields::from(vec![
7428 key_field.as_ref().clone(),
7429 val_field.as_ref().clone(),
7430 ]),
7431 vec![Arc::new(keys) as ArrayRef, vals],
7432 None,
7433 );
7434 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0]));
7435 Arc::new(MapArray::new(
7436 entry_field.clone(),
7437 offsets,
7438 entries,
7439 None,
7440 *is_sorted,
7441 ))
7442 }
7443 other => panic!("empty_child_for: unhandled type {other:?}"),
7444 }
7445 }
7446 let children: Vec<ArrayRef> = fields
7447 .iter()
7448 .map(|(_, f)| provide(f).unwrap_or_else(|| empty_child_for(f.data_type())))
7449 .collect();
7450 Arc::new(
7451 UnionArray::try_new(
7452 fields.clone(),
7453 ScalarBuffer::<i8>::from(type_ids),
7454 Some(ScalarBuffer::<i32>::from(offsets)),
7455 children,
7456 )
7457 .unwrap(),
7458 ) as ArrayRef
7459 }
7460
7461 #[inline]
7462 fn uuid16_from_str(s: &str) -> [u8; 16] {
7463 let mut out = [0u8; 16];
7464 let mut idx = 0usize;
7465 let mut hi: Option<u8> = None;
7466 for ch in s.chars() {
7467 if ch == '-' {
7468 continue;
7469 }
7470 let v = ch.to_digit(16).expect("invalid hex digit in UUID") as u8;
7471 if let Some(h) = hi {
7472 out[idx] = (h << 4) | v;
7473 idx += 1;
7474 hi = None;
7475 } else {
7476 hi = Some(v);
7477 }
7478 }
7479 assert_eq!(idx, 16, "UUID must decode to 16 bytes");
7480 out
7481 }
7482 let date_a: i32 = 19_000; let time_ms_a: i32 = 12 * 3_600_000 + 34 * 60_000 + 56_000 + 789;
7484 let time_us_eod: i64 = 86_400_000_000 - 1;
7485 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;
7487 let dur_small = IntervalMonthDayNanoType::make_value(1, 2, 3_000_000_000);
7488 let dur_zero = IntervalMonthDayNanoType::make_value(0, 0, 0);
7489 let dur_large =
7490 IntervalMonthDayNanoType::make_value(12, 31, ((86_400_000 - 1) as i64) * 1_000_000);
7491 let dur_2years = IntervalMonthDayNanoType::make_value(24, 0, 0);
7492 let uuid1 = uuid16_from_str("fe7bc30b-4ce8-4c5e-b67c-2234a2d38e66");
7493 let uuid2 = uuid16_from_str("0826cc06-d2e3-4599-b4ad-af5fa6905cdb");
7494
7495 #[inline]
7496 fn push_like(
7497 reader_schema: &arrow_schema::Schema,
7498 name: &str,
7499 arr: ArrayRef,
7500 fields: &mut Vec<FieldRef>,
7501 cols: &mut Vec<ArrayRef>,
7502 ) {
7503 let src = reader_schema
7504 .field_with_name(name)
7505 .unwrap_or_else(|_| panic!("source schema missing field '{name}'"));
7506 let mut f = Field::new(name, arr.data_type().clone(), src.is_nullable());
7507 let md = src.metadata();
7508 if !md.is_empty() {
7509 f = f.with_metadata(md.clone());
7510 }
7511 fields.push(Arc::new(f));
7512 cols.push(arr);
7513 }
7514
7515 let mut fields: Vec<FieldRef> = Vec::new();
7516 let mut columns: Vec<ArrayRef> = Vec::new();
7517 push_like(
7518 schema.as_ref(),
7519 "id",
7520 Arc::new(Int64Array::from(vec![1, 2, 3, 4])) as ArrayRef,
7521 &mut fields,
7522 &mut columns,
7523 );
7524 push_like(
7525 schema.as_ref(),
7526 "flag",
7527 Arc::new(BooleanArray::from(vec![true, false, true, false])) as ArrayRef,
7528 &mut fields,
7529 &mut columns,
7530 );
7531 push_like(
7532 schema.as_ref(),
7533 "ratio_f32",
7534 Arc::new(Float32Array::from(vec![1.25f32, -0.0, 3.5, 9.75])) as ArrayRef,
7535 &mut fields,
7536 &mut columns,
7537 );
7538 push_like(
7539 schema.as_ref(),
7540 "ratio_f64",
7541 Arc::new(Float64Array::from(vec![2.5f64, -1.0, 7.0, -2.25])) as ArrayRef,
7542 &mut fields,
7543 &mut columns,
7544 );
7545 push_like(
7546 schema.as_ref(),
7547 "count_i32",
7548 Arc::new(Int32Array::from(vec![7, -1, 0, 123])) as ArrayRef,
7549 &mut fields,
7550 &mut columns,
7551 );
7552 push_like(
7553 schema.as_ref(),
7554 "count_i64",
7555 Arc::new(Int64Array::from(vec![
7556 7_000_000_000i64,
7557 -2,
7558 0,
7559 -9_876_543_210i64,
7560 ])) as ArrayRef,
7561 &mut fields,
7562 &mut columns,
7563 );
7564 push_like(
7565 schema.as_ref(),
7566 "opt_i32_nullfirst",
7567 Arc::new(Int32Array::from(vec![None, Some(42), None, Some(0)])) as ArrayRef,
7568 &mut fields,
7569 &mut columns,
7570 );
7571 push_like(
7572 schema.as_ref(),
7573 "opt_str_nullsecond",
7574 Arc::new(StringArray::from(vec![
7575 Some("alpha"),
7576 None,
7577 Some("s3"),
7578 Some(""),
7579 ])) as ArrayRef,
7580 &mut fields,
7581 &mut columns,
7582 );
7583 {
7584 let uf = match schema
7585 .field_with_name("tri_union_prim")
7586 .unwrap()
7587 .data_type()
7588 {
7589 DataType::Union(f, UnionMode::Dense) => f.clone(),
7590 other => panic!("tri_union_prim should be dense union, got {other:?}"),
7591 };
7592 let tid_i = tid_by_name(&uf, "int");
7593 let tid_s = tid_by_name(&uf, "string");
7594 let tid_b = tid_by_name(&uf, "boolean");
7595 let tids = vec![tid_i, tid_s, tid_b, tid_s];
7596 let offs = vec![0, 0, 0, 1];
7597 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
7598 DataType::Int32 => Some(Arc::new(Int32Array::from(vec![0])) as ArrayRef),
7599 DataType::Utf8 => Some(Arc::new(StringArray::from(vec!["hi", ""])) as ArrayRef),
7600 DataType::Boolean => Some(Arc::new(BooleanArray::from(vec![true])) as ArrayRef),
7601 _ => None,
7602 });
7603 push_like(
7604 schema.as_ref(),
7605 "tri_union_prim",
7606 arr,
7607 &mut fields,
7608 &mut columns,
7609 );
7610 }
7611
7612 push_like(
7613 schema.as_ref(),
7614 "str_utf8",
7615 Arc::new(StringArray::from(vec!["hello", "", "world", "✓ unicode"])) as ArrayRef,
7616 &mut fields,
7617 &mut columns,
7618 );
7619 push_like(
7620 schema.as_ref(),
7621 "raw_bytes",
7622 Arc::new(BinaryArray::from(vec![
7623 b"\x00\x01".as_ref(),
7624 b"".as_ref(),
7625 b"\xFF\x00".as_ref(),
7626 b"\x10\x20\x30\x40".as_ref(),
7627 ])) as ArrayRef,
7628 &mut fields,
7629 &mut columns,
7630 );
7631 {
7632 let it = [
7633 Some(*b"0123456789ABCDEF"),
7634 Some([0u8; 16]),
7635 Some(*b"ABCDEFGHIJKLMNOP"),
7636 Some([0xAA; 16]),
7637 ]
7638 .into_iter();
7639 let arr =
7640 Arc::new(FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 16).unwrap())
7641 as ArrayRef;
7642 push_like(
7643 schema.as_ref(),
7644 "fx16_plain",
7645 arr,
7646 &mut fields,
7647 &mut columns,
7648 );
7649 }
7650 {
7651 #[cfg(feature = "small_decimals")]
7652 let dec10_2 = Arc::new(
7653 Decimal64Array::from_iter_values([123456i64, -1, 0, 9_999_999_999i64])
7654 .with_precision_and_scale(10, 2)
7655 .unwrap(),
7656 ) as ArrayRef;
7657 #[cfg(not(feature = "small_decimals"))]
7658 let dec10_2 = Arc::new(
7659 Decimal128Array::from_iter_values([123456i128, -1, 0, 9_999_999_999i128])
7660 .with_precision_and_scale(10, 2)
7661 .unwrap(),
7662 ) as ArrayRef;
7663 push_like(
7664 schema.as_ref(),
7665 "dec_bytes_s10_2",
7666 dec10_2,
7667 &mut fields,
7668 &mut columns,
7669 );
7670 }
7671 {
7672 #[cfg(feature = "small_decimals")]
7673 let dec20_4 = Arc::new(
7674 Decimal128Array::from_iter_values([1_234_567_891_234i128, -420_000i128, 0, -1i128])
7675 .with_precision_and_scale(20, 4)
7676 .unwrap(),
7677 ) as ArrayRef;
7678 #[cfg(not(feature = "small_decimals"))]
7679 let dec20_4 = Arc::new(
7680 Decimal128Array::from_iter_values([1_234_567_891_234i128, -420_000i128, 0, -1i128])
7681 .with_precision_and_scale(20, 4)
7682 .unwrap(),
7683 ) as ArrayRef;
7684 push_like(
7685 schema.as_ref(),
7686 "dec_fix_s20_4",
7687 dec20_4,
7688 &mut fields,
7689 &mut columns,
7690 );
7691 }
7692 {
7693 let it = [Some(uuid1), Some(uuid2), Some(uuid1), Some(uuid2)].into_iter();
7694 let arr =
7695 Arc::new(FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 16).unwrap())
7696 as ArrayRef;
7697 push_like(schema.as_ref(), "uuid_str", arr, &mut fields, &mut columns);
7698 }
7699 push_like(
7700 schema.as_ref(),
7701 "d_date",
7702 Arc::new(Date32Array::from(vec![date_a, 0, 1, 365])) as ArrayRef,
7703 &mut fields,
7704 &mut columns,
7705 );
7706 push_like(
7707 schema.as_ref(),
7708 "t_millis",
7709 Arc::new(Time32MillisecondArray::from(vec![
7710 time_ms_a,
7711 0,
7712 1,
7713 86_400_000 - 1,
7714 ])) as ArrayRef,
7715 &mut fields,
7716 &mut columns,
7717 );
7718 push_like(
7719 schema.as_ref(),
7720 "t_micros",
7721 Arc::new(Time64MicrosecondArray::from(vec![
7722 time_us_eod,
7723 0,
7724 1,
7725 1_000_000,
7726 ])) as ArrayRef,
7727 &mut fields,
7728 &mut columns,
7729 );
7730 {
7731 let a = TimestampMillisecondArray::from(vec![
7732 ts_ms_2024_01_01,
7733 -1,
7734 ts_ms_2024_01_01 + 123,
7735 0,
7736 ])
7737 .with_timezone("+00:00");
7738 push_like(
7739 schema.as_ref(),
7740 "ts_millis_utc",
7741 Arc::new(a) as ArrayRef,
7742 &mut fields,
7743 &mut columns,
7744 );
7745 }
7746 {
7747 let a = TimestampMicrosecondArray::from(vec![
7748 ts_us_2024_01_01,
7749 1,
7750 ts_us_2024_01_01 + 456,
7751 0,
7752 ])
7753 .with_timezone("+00:00");
7754 push_like(
7755 schema.as_ref(),
7756 "ts_micros_utc",
7757 Arc::new(a) as ArrayRef,
7758 &mut fields,
7759 &mut columns,
7760 );
7761 }
7762 push_like(
7763 schema.as_ref(),
7764 "ts_millis_local",
7765 Arc::new(TimestampMillisecondArray::from(vec![
7766 ts_ms_2024_01_01 + 86_400_000,
7767 0,
7768 ts_ms_2024_01_01 + 789,
7769 123_456_789,
7770 ])) as ArrayRef,
7771 &mut fields,
7772 &mut columns,
7773 );
7774 push_like(
7775 schema.as_ref(),
7776 "ts_micros_local",
7777 Arc::new(TimestampMicrosecondArray::from(vec![
7778 ts_us_2024_01_01 + 123_456,
7779 0,
7780 ts_us_2024_01_01 + 101_112,
7781 987_654_321,
7782 ])) as ArrayRef,
7783 &mut fields,
7784 &mut columns,
7785 );
7786 {
7787 let v = vec![dur_small, dur_zero, dur_large, dur_2years];
7788 push_like(
7789 schema.as_ref(),
7790 "interval_mdn",
7791 Arc::new(IntervalMonthDayNanoArray::from(v)) as ArrayRef,
7792 &mut fields,
7793 &mut columns,
7794 );
7795 }
7796 {
7797 let keys = Int32Array::from(vec![1, 2, 3, 0]); let values = Arc::new(StringArray::from(vec![
7799 "UNKNOWN",
7800 "NEW",
7801 "PROCESSING",
7802 "DONE",
7803 ])) as ArrayRef;
7804 let dict = DictionaryArray::<Int32Type>::try_new(keys, values).unwrap();
7805 push_like(
7806 schema.as_ref(),
7807 "status",
7808 Arc::new(dict) as ArrayRef,
7809 &mut fields,
7810 &mut columns,
7811 );
7812 }
7813 {
7814 let list_field = match schema.field_with_name("arr_union").unwrap().data_type() {
7815 DataType::List(f) => f.clone(),
7816 other => panic!("arr_union should be List, got {other:?}"),
7817 };
7818 let uf = match list_field.data_type() {
7819 DataType::Union(f, UnionMode::Dense) => f.clone(),
7820 other => panic!("arr_union item should be union, got {other:?}"),
7821 };
7822 let tid_l = tid_by_name(&uf, "long");
7823 let tid_s = tid_by_name(&uf, "string");
7824 let tid_n = tid_by_name(&uf, "null");
7825 let type_ids = vec![
7826 tid_l, tid_s, tid_n, tid_l, tid_n, tid_s, tid_l, tid_l, tid_s, tid_n, tid_l,
7827 ];
7828 let offsets = vec![0, 0, 0, 1, 1, 1, 2, 3, 2, 2, 4];
7829 let values = mk_dense_union(&uf, type_ids, offsets, |f| match f.data_type() {
7830 DataType::Int64 => {
7831 Some(Arc::new(Int64Array::from(vec![1i64, -3, 0, -1, 0])) as ArrayRef)
7832 }
7833 DataType::Utf8 => {
7834 Some(Arc::new(StringArray::from(vec!["x", "z", "end"])) as ArrayRef)
7835 }
7836 DataType::Null => Some(Arc::new(NullArray::new(3)) as ArrayRef),
7837 _ => None,
7838 });
7839 let list_offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 4, 7, 8, 11]));
7840 let arr = Arc::new(ListArray::try_new(list_field, list_offsets, values, None).unwrap())
7841 as ArrayRef;
7842 push_like(schema.as_ref(), "arr_union", arr, &mut fields, &mut columns);
7843 }
7844 {
7845 let (entry_field, entries_fields, uf, is_sorted) =
7846 match schema.field_with_name("map_union").unwrap().data_type() {
7847 DataType::Map(entry_field, is_sorted) => {
7848 let fs = match entry_field.data_type() {
7849 DataType::Struct(fs) => fs.clone(),
7850 other => panic!("map entries must be struct, got {other:?}"),
7851 };
7852 let val_f = fs[1].clone();
7853 let uf = match val_f.data_type() {
7854 DataType::Union(f, UnionMode::Dense) => f.clone(),
7855 other => panic!("map value must be union, got {other:?}"),
7856 };
7857 (entry_field.clone(), fs, uf, *is_sorted)
7858 }
7859 other => panic!("map_union should be Map, got {other:?}"),
7860 };
7861 let keys = StringArray::from(vec!["a", "b", "c", "neg", "pi", "ok"]);
7862 let moff = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 4, 4, 6]));
7863 let tid_null = tid_by_name(&uf, "null");
7864 let tid_d = tid_by_name(&uf, "double");
7865 let tid_s = tid_by_name(&uf, "string");
7866 let type_ids = vec![tid_d, tid_null, tid_s, tid_d, tid_d, tid_s];
7867 let offsets = vec![0, 0, 0, 1, 2, 1];
7868 let pi_5dp = (std::f64::consts::PI * 100_000.0).trunc() / 100_000.0;
7869 let vals = mk_dense_union(&uf, type_ids, offsets, |f| match f.data_type() {
7870 DataType::Float64 => {
7871 Some(Arc::new(Float64Array::from(vec![1.5f64, -0.5, pi_5dp])) as ArrayRef)
7872 }
7873 DataType::Utf8 => {
7874 Some(Arc::new(StringArray::from(vec!["yes", "true"])) as ArrayRef)
7875 }
7876 DataType::Null => Some(Arc::new(NullArray::new(2)) as ArrayRef),
7877 _ => None,
7878 });
7879 let entries = StructArray::new(
7880 entries_fields.clone(),
7881 vec![Arc::new(keys) as ArrayRef, vals],
7882 None,
7883 );
7884 let map =
7885 Arc::new(MapArray::new(entry_field, moff, entries, None, is_sorted)) as ArrayRef;
7886 push_like(schema.as_ref(), "map_union", map, &mut fields, &mut columns);
7887 }
7888 {
7889 let fs = match schema.field_with_name("address").unwrap().data_type() {
7890 DataType::Struct(fs) => fs.clone(),
7891 other => panic!("address should be Struct, got {other:?}"),
7892 };
7893 let street = Arc::new(StringArray::from(vec![
7894 "100 Main",
7895 "",
7896 "42 Galaxy Way",
7897 "End Ave",
7898 ])) as ArrayRef;
7899 let zip = Arc::new(Int32Array::from(vec![12345, 0, 42424, 1])) as ArrayRef;
7900 let country = Arc::new(StringArray::from(vec!["US", "CA", "US", "GB"])) as ArrayRef;
7901 let arr = Arc::new(StructArray::new(fs, vec![street, zip, country], None)) as ArrayRef;
7902 push_like(schema.as_ref(), "address", arr, &mut fields, &mut columns);
7903 }
7904 {
7905 let fs = match schema.field_with_name("maybe_auth").unwrap().data_type() {
7906 DataType::Struct(fs) => fs.clone(),
7907 other => panic!("maybe_auth should be Struct, got {other:?}"),
7908 };
7909 let user =
7910 Arc::new(StringArray::from(vec!["alice", "bob", "carol", "dave"])) as ArrayRef;
7911 let token_values: Vec<Option<&[u8]>> = vec![
7912 None, Some(b"\x01\x02\x03".as_ref()), None, Some(b"".as_ref()), ];
7917 let token = Arc::new(BinaryArray::from(token_values)) as ArrayRef;
7918 let arr = Arc::new(StructArray::new(fs, vec![user, token], None)) as ArrayRef;
7919 push_like(
7920 schema.as_ref(),
7921 "maybe_auth",
7922 arr,
7923 &mut fields,
7924 &mut columns,
7925 );
7926 }
7927 {
7928 let uf = match schema
7929 .field_with_name("union_enum_record_array_map")
7930 .unwrap()
7931 .data_type()
7932 {
7933 DataType::Union(f, UnionMode::Dense) => f.clone(),
7934 other => panic!("union_enum_record_array_map should be union, got {other:?}"),
7935 };
7936 let mut tid_enum: Option<i8> = None;
7937 let mut tid_rec_a: Option<i8> = None;
7938 let mut tid_array: Option<i8> = None;
7939 let mut tid_map: Option<i8> = None;
7940 let mut map_entry_field: Option<FieldRef> = None;
7941 let mut map_sorted: bool = false;
7942 for (tid, f) in uf.iter() {
7943 match f.data_type() {
7944 DataType::Dictionary(_, _) => tid_enum = Some(tid),
7945 DataType::Struct(childs)
7946 if childs.len() == 2
7947 && childs[0].name() == "a"
7948 && childs[1].name() == "b" =>
7949 {
7950 tid_rec_a = Some(tid)
7951 }
7952 DataType::List(item) if matches!(item.data_type(), DataType::Int64) => {
7953 tid_array = Some(tid)
7954 }
7955 DataType::Map(ef, is_sorted) => {
7956 tid_map = Some(tid);
7957 map_entry_field = Some(ef.clone());
7958 map_sorted = *is_sorted;
7959 }
7960 _ => {}
7961 }
7962 }
7963 let (tid_enum, tid_rec_a, tid_array, tid_map) = (
7964 tid_enum.unwrap(),
7965 tid_rec_a.unwrap(),
7966 tid_array.unwrap(),
7967 tid_map.unwrap(),
7968 );
7969 let tids = vec![tid_enum, tid_rec_a, tid_array, tid_map];
7970 let offs = vec![0, 0, 0, 0];
7971 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
7972 DataType::Dictionary(_, _) => {
7973 let keys = Int32Array::from(vec![0i32]);
7974 let values =
7975 Arc::new(StringArray::from(vec!["RED", "GREEN", "BLUE"])) as ArrayRef;
7976 Some(
7977 Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap())
7978 as ArrayRef,
7979 )
7980 }
7981 DataType::Struct(fs)
7982 if fs.len() == 2 && fs[0].name() == "a" && fs[1].name() == "b" =>
7983 {
7984 let a = Int32Array::from(vec![7]);
7985 let b = StringArray::from(vec!["rec"]);
7986 Some(Arc::new(StructArray::new(
7987 fs.clone(),
7988 vec![Arc::new(a), Arc::new(b)],
7989 None,
7990 )) as ArrayRef)
7991 }
7992 DataType::List(field) => {
7993 let values = Int64Array::from(vec![1i64, 2, 3]);
7994 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3]));
7995 Some(Arc::new(
7996 ListArray::try_new(field.clone(), offsets, Arc::new(values), None).unwrap(),
7997 ) as ArrayRef)
7998 }
7999 DataType::Map(_, _) => {
8000 let entry_field = map_entry_field.clone().unwrap();
8001 let (key_field, val_field) = match entry_field.data_type() {
8002 DataType::Struct(fs) => (fs[0].clone(), fs[1].clone()),
8003 _ => unreachable!(),
8004 };
8005 let keys = StringArray::from(vec!["k"]);
8006 let vals = StringArray::from(vec!["v"]);
8007 let entries = StructArray::new(
8008 Fields::from(vec![key_field.as_ref().clone(), val_field.as_ref().clone()]),
8009 vec![Arc::new(keys) as ArrayRef, Arc::new(vals) as ArrayRef],
8010 None,
8011 );
8012 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 1]));
8013 Some(Arc::new(MapArray::new(
8014 entry_field.clone(),
8015 offsets,
8016 entries,
8017 None,
8018 map_sorted,
8019 )) as ArrayRef)
8020 }
8021 _ => None,
8022 });
8023 push_like(
8024 schema.as_ref(),
8025 "union_enum_record_array_map",
8026 arr,
8027 &mut fields,
8028 &mut columns,
8029 );
8030 }
8031 {
8032 let uf = match schema
8033 .field_with_name("union_date_or_fixed4")
8034 .unwrap()
8035 .data_type()
8036 {
8037 DataType::Union(f, UnionMode::Dense) => f.clone(),
8038 other => panic!("union_date_or_fixed4 should be union, got {other:?}"),
8039 };
8040 let tid_date = tid_by_dt(&uf, |dt| matches!(dt, DataType::Date32));
8041 let tid_fx4 = tid_by_dt(&uf, |dt| matches!(dt, DataType::FixedSizeBinary(4)));
8042 let tids = vec![tid_date, tid_fx4, tid_date, tid_fx4];
8043 let offs = vec![0, 0, 1, 1];
8044 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
8045 DataType::Date32 => Some(Arc::new(Date32Array::from(vec![date_a, 0])) as ArrayRef),
8046 DataType::FixedSizeBinary(4) => {
8047 let it = [Some(*b"\x00\x11\x22\x33"), Some(*b"ABCD")].into_iter();
8048 Some(Arc::new(
8049 FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 4).unwrap(),
8050 ) as ArrayRef)
8051 }
8052 _ => None,
8053 });
8054 push_like(
8055 schema.as_ref(),
8056 "union_date_or_fixed4",
8057 arr,
8058 &mut fields,
8059 &mut columns,
8060 );
8061 }
8062 {
8063 let uf = match schema
8064 .field_with_name("union_interval_or_string")
8065 .unwrap()
8066 .data_type()
8067 {
8068 DataType::Union(f, UnionMode::Dense) => f.clone(),
8069 other => panic!("union_interval_or_string should be union, got {other:?}"),
8070 };
8071 let tid_dur = tid_by_dt(&uf, |dt| {
8072 matches!(dt, DataType::Interval(IntervalUnit::MonthDayNano))
8073 });
8074 let tid_str = tid_by_dt(&uf, |dt| matches!(dt, DataType::Utf8));
8075 let tids = vec![tid_dur, tid_str, tid_dur, tid_str];
8076 let offs = vec![0, 0, 1, 1];
8077 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
8078 DataType::Interval(IntervalUnit::MonthDayNano) => Some(Arc::new(
8079 IntervalMonthDayNanoArray::from(vec![dur_small, dur_large]),
8080 )
8081 as ArrayRef),
8082 DataType::Utf8 => Some(Arc::new(StringArray::from(vec![
8083 "duration-as-text",
8084 "iso-8601-period-P1Y",
8085 ])) as ArrayRef),
8086 _ => None,
8087 });
8088 push_like(
8089 schema.as_ref(),
8090 "union_interval_or_string",
8091 arr,
8092 &mut fields,
8093 &mut columns,
8094 );
8095 }
8096 {
8097 let uf = match schema
8098 .field_with_name("union_uuid_or_fixed10")
8099 .unwrap()
8100 .data_type()
8101 {
8102 DataType::Union(f, UnionMode::Dense) => f.clone(),
8103 other => panic!("union_uuid_or_fixed10 should be union, got {other:?}"),
8104 };
8105 let tid_uuid = tid_by_dt(&uf, |dt| matches!(dt, DataType::FixedSizeBinary(16)));
8106 let tid_fx10 = tid_by_dt(&uf, |dt| matches!(dt, DataType::FixedSizeBinary(10)));
8107 let tids = vec![tid_uuid, tid_fx10, tid_uuid, tid_fx10];
8108 let offs = vec![0, 0, 1, 1];
8109 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
8110 DataType::FixedSizeBinary(16) => {
8111 let it = [Some(uuid1), Some(uuid2)].into_iter();
8112 Some(Arc::new(
8113 FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 16).unwrap(),
8114 ) as ArrayRef)
8115 }
8116 DataType::FixedSizeBinary(10) => {
8117 let fx10_a = [0xAAu8; 10];
8118 let fx10_b = [0x00u8, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99];
8119 let it = [Some(fx10_a), Some(fx10_b)].into_iter();
8120 Some(Arc::new(
8121 FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 10).unwrap(),
8122 ) as ArrayRef)
8123 }
8124 _ => None,
8125 });
8126 push_like(
8127 schema.as_ref(),
8128 "union_uuid_or_fixed10",
8129 arr,
8130 &mut fields,
8131 &mut columns,
8132 );
8133 }
8134 {
8135 let list_field = match schema
8136 .field_with_name("array_records_with_union")
8137 .unwrap()
8138 .data_type()
8139 {
8140 DataType::List(f) => f.clone(),
8141 other => panic!("array_records_with_union should be List, got {other:?}"),
8142 };
8143 let kv_fields = match list_field.data_type() {
8144 DataType::Struct(fs) => fs.clone(),
8145 other => panic!("array_records_with_union items must be Struct, got {other:?}"),
8146 };
8147 let val_field = kv_fields
8148 .iter()
8149 .find(|f| f.name() == "val")
8150 .unwrap()
8151 .clone();
8152 let uf = match val_field.data_type() {
8153 DataType::Union(f, UnionMode::Dense) => f.clone(),
8154 other => panic!("KV.val should be union, got {other:?}"),
8155 };
8156 let keys = Arc::new(StringArray::from(vec!["k1", "k2", "k", "k3", "x"])) as ArrayRef;
8157 let tid_null = tid_by_name(&uf, "null");
8158 let tid_i = tid_by_name(&uf, "int");
8159 let tid_l = tid_by_name(&uf, "long");
8160 let type_ids = vec![tid_i, tid_null, tid_l, tid_null, tid_i];
8161 let offsets = vec![0, 0, 0, 1, 1];
8162 let vals = mk_dense_union(&uf, type_ids, offsets, |f| match f.data_type() {
8163 DataType::Int32 => Some(Arc::new(Int32Array::from(vec![5, -5])) as ArrayRef),
8164 DataType::Int64 => Some(Arc::new(Int64Array::from(vec![99i64])) as ArrayRef),
8165 DataType::Null => Some(Arc::new(NullArray::new(2)) as ArrayRef),
8166 _ => None,
8167 });
8168 let values_struct =
8169 Arc::new(StructArray::new(kv_fields.clone(), vec![keys, vals], None)) as ArrayRef;
8170 let list_offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2, 3, 4, 5]));
8171 let arr = Arc::new(
8172 ListArray::try_new(list_field, list_offsets, values_struct, None).unwrap(),
8173 ) as ArrayRef;
8174 push_like(
8175 schema.as_ref(),
8176 "array_records_with_union",
8177 arr,
8178 &mut fields,
8179 &mut columns,
8180 );
8181 }
8182 {
8183 let uf = match schema
8184 .field_with_name("union_map_or_array_int")
8185 .unwrap()
8186 .data_type()
8187 {
8188 DataType::Union(f, UnionMode::Dense) => f.clone(),
8189 other => panic!("union_map_or_array_int should be union, got {other:?}"),
8190 };
8191 let tid_map = tid_by_dt(&uf, |dt| matches!(dt, DataType::Map(_, _)));
8192 let tid_list = tid_by_dt(&uf, |dt| matches!(dt, DataType::List(_)));
8193 let map_child: ArrayRef = {
8194 let (entry_field, is_sorted) = match uf
8195 .iter()
8196 .find(|(tid, _)| *tid == tid_map)
8197 .unwrap()
8198 .1
8199 .data_type()
8200 {
8201 DataType::Map(ef, is_sorted) => (ef.clone(), *is_sorted),
8202 _ => unreachable!(),
8203 };
8204 let (key_field, val_field) = match entry_field.data_type() {
8205 DataType::Struct(fs) => (fs[0].clone(), fs[1].clone()),
8206 _ => unreachable!(),
8207 };
8208 let keys = StringArray::from(vec!["x", "y", "only"]);
8209 let vals = Int32Array::from(vec![1, 2, 10]);
8210 let entries = StructArray::new(
8211 Fields::from(vec![key_field.as_ref().clone(), val_field.as_ref().clone()]),
8212 vec![Arc::new(keys) as ArrayRef, Arc::new(vals) as ArrayRef],
8213 None,
8214 );
8215 let moff = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2, 3]));
8216 Arc::new(MapArray::new(entry_field, moff, entries, None, is_sorted)) as ArrayRef
8217 };
8218 let list_child: ArrayRef = {
8219 let list_field = match uf
8220 .iter()
8221 .find(|(tid, _)| *tid == tid_list)
8222 .unwrap()
8223 .1
8224 .data_type()
8225 {
8226 DataType::List(f) => f.clone(),
8227 _ => unreachable!(),
8228 };
8229 let values = Int32Array::from(vec![1, 2, 3, 0]);
8230 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 4]));
8231 Arc::new(ListArray::try_new(list_field, offsets, Arc::new(values), None).unwrap())
8232 as ArrayRef
8233 };
8234 let tids = vec![tid_map, tid_list, tid_map, tid_list];
8235 let offs = vec![0, 0, 1, 1];
8236 let arr = mk_dense_union(&uf, tids, offs, |f| match f.data_type() {
8237 DataType::Map(_, _) => Some(map_child.clone()),
8238 DataType::List(_) => Some(list_child.clone()),
8239 _ => None,
8240 });
8241 push_like(
8242 schema.as_ref(),
8243 "union_map_or_array_int",
8244 arr,
8245 &mut fields,
8246 &mut columns,
8247 );
8248 }
8249 push_like(
8250 schema.as_ref(),
8251 "renamed_with_default",
8252 Arc::new(Int32Array::from(vec![100, 42, 7, 42])) as ArrayRef,
8253 &mut fields,
8254 &mut columns,
8255 );
8256 {
8257 let fs = match schema.field_with_name("person").unwrap().data_type() {
8258 DataType::Struct(fs) => fs.clone(),
8259 other => panic!("person should be Struct, got {other:?}"),
8260 };
8261 let name =
8262 Arc::new(StringArray::from(vec!["Alice", "Bob", "Carol", "Dave"])) as ArrayRef;
8263 let age = Arc::new(Int32Array::from(vec![30, 0, 25, 41])) as ArrayRef;
8264 let arr = Arc::new(StructArray::new(fs, vec![name, age], None)) as ArrayRef;
8265 push_like(schema.as_ref(), "person", arr, &mut fields, &mut columns);
8266 }
8267 let expected =
8268 RecordBatch::try_new(Arc::new(Schema::new(Fields::from(fields))), columns).unwrap();
8269 assert_eq!(
8270 expected, batch,
8271 "entire RecordBatch mismatch (schema, all columns, all rows)"
8272 );
8273 }
8274 #[test]
8275 fn comprehensive_e2e_resolution_test() {
8276 use serde_json::Value;
8277 use std::collections::HashMap;
8278
8279 fn make_comprehensive_reader_schema(path: &str) -> AvroSchema {
8292 fn set_type_string(f: &mut Value, new_ty: &str) {
8293 if let Some(ty) = f.get_mut("type") {
8294 match ty {
8295 Value::String(_) | Value::Object(_) => {
8296 *ty = Value::String(new_ty.to_string());
8297 }
8298 Value::Array(arr) => {
8299 for b in arr.iter_mut() {
8300 match b {
8301 Value::String(s) if s != "null" => {
8302 *b = Value::String(new_ty.to_string());
8303 break;
8304 }
8305 Value::Object(_) => {
8306 *b = Value::String(new_ty.to_string());
8307 break;
8308 }
8309 _ => {}
8310 }
8311 }
8312 }
8313 _ => {}
8314 }
8315 }
8316 }
8317 fn reverse_union_array(f: &mut Value) {
8318 if let Some(arr) = f.get_mut("type").and_then(|t| t.as_array_mut()) {
8319 arr.reverse();
8320 }
8321 }
8322 fn reverse_items_union(f: &mut Value) {
8323 if let Some(obj) = f.get_mut("type").and_then(|t| t.as_object_mut())
8324 && let Some(items) = obj.get_mut("items").and_then(|v| v.as_array_mut())
8325 {
8326 items.reverse();
8327 }
8328 }
8329 fn reverse_map_values_union(f: &mut Value) {
8330 if let Some(obj) = f.get_mut("type").and_then(|t| t.as_object_mut())
8331 && let Some(values) = obj.get_mut("values").and_then(|v| v.as_array_mut())
8332 {
8333 values.reverse();
8334 }
8335 }
8336 fn reverse_nested_union_in_record(f: &mut Value, field_name: &str) {
8337 if let Some(obj) = f.get_mut("type").and_then(|t| t.as_object_mut())
8338 && let Some(fields) = obj.get_mut("fields").and_then(|v| v.as_array_mut())
8339 {
8340 for ff in fields.iter_mut() {
8341 if ff.get("name").and_then(|n| n.as_str()) == Some(field_name)
8342 && let Some(ty) = ff.get_mut("type")
8343 && let Some(arr) = ty.as_array_mut()
8344 {
8345 arr.reverse();
8346 }
8347 }
8348 }
8349 }
8350 fn rename_nested_field_with_alias(f: &mut Value, old: &str, new: &str) {
8351 if let Some(obj) = f.get_mut("type").and_then(|t| t.as_object_mut())
8352 && let Some(fields) = obj.get_mut("fields").and_then(|v| v.as_array_mut())
8353 {
8354 for ff in fields.iter_mut() {
8355 if ff.get("name").and_then(|n| n.as_str()) == Some(old) {
8356 ff["name"] = Value::String(new.to_string());
8357 ff["aliases"] = Value::Array(vec![Value::String(old.to_string())]);
8358 }
8359 }
8360 }
8361 }
8362 let mut root = load_writer_schema_json(path);
8363 assert_eq!(root["type"], "record", "writer schema must be a record");
8364 let fields = root
8365 .get_mut("fields")
8366 .and_then(|f| f.as_array_mut())
8367 .expect("record has fields");
8368 for f in fields.iter_mut() {
8369 let Some(name) = f.get("name").and_then(|n| n.as_str()) else {
8370 continue;
8371 };
8372 match name {
8373 "id" => {
8375 f["name"] = Value::String("identifier".into());
8376 f["aliases"] = Value::Array(vec![Value::String("id".into())]);
8377 }
8378 "renamed_with_default" => {
8379 f["name"] = Value::String("old_count".into());
8380 f["aliases"] =
8381 Value::Array(vec![Value::String("renamed_with_default".into())]);
8382 }
8383 "count_i32" => set_type_string(f, "long"),
8385 "ratio_f32" => set_type_string(f, "double"),
8386 "opt_str_nullsecond" => reverse_union_array(f),
8388 "union_enum_record_array_map" => reverse_union_array(f),
8389 "union_date_or_fixed4" => reverse_union_array(f),
8390 "union_interval_or_string" => reverse_union_array(f),
8391 "union_uuid_or_fixed10" => reverse_union_array(f),
8392 "union_map_or_array_int" => reverse_union_array(f),
8393 "maybe_auth" => reverse_nested_union_in_record(f, "token"),
8394 "arr_union" => reverse_items_union(f),
8396 "map_union" => reverse_map_values_union(f),
8397 "address" => rename_nested_field_with_alias(f, "street", "street_name"),
8399 "person" => {
8401 if let Some(tobj) = f.get_mut("type").and_then(|t| t.as_object_mut()) {
8402 tobj.insert("name".to_string(), Value::String("Person".into()));
8403 tobj.insert(
8404 "namespace".to_string(),
8405 Value::String("com.example".into()),
8406 );
8407 tobj.insert(
8408 "aliases".into(),
8409 Value::Array(vec![
8410 Value::String("PersonV2".into()),
8411 Value::String("com.example.v2.PersonV2".into()),
8412 ]),
8413 );
8414 }
8415 }
8416 _ => {}
8417 }
8418 }
8419 fields.reverse();
8420 AvroSchema::new(root.to_string())
8421 }
8422
8423 let path = "test/data/comprehensive_e2e.avro";
8424 let reader_schema = make_comprehensive_reader_schema(path);
8425 let batch = read_alltypes_with_reader_schema(path, reader_schema.clone());
8426
8427 const UUID_EXT_KEY: &str = "ARROW:extension:name";
8428 const UUID_LOGICAL_KEY: &str = "logicalType";
8429
8430 let uuid_md_top: Option<arrow_schema::Metadata> = batch
8431 .schema()
8432 .field_with_name("uuid_str")
8433 .ok()
8434 .and_then(|f| {
8435 let md = f.metadata();
8436 let has_ext = md.get(UUID_EXT_KEY).is_some();
8437 let is_uuid_logical = md
8438 .get(UUID_LOGICAL_KEY)
8439 .map(|v| v.trim_matches('"') == "uuid")
8440 .unwrap_or(false);
8441 if has_ext || is_uuid_logical {
8442 Some(md.clone())
8443 } else {
8444 None
8445 }
8446 });
8447
8448 let uuid_md_union: Option<arrow_schema::Metadata> = batch
8449 .schema()
8450 .field_with_name("union_uuid_or_fixed10")
8451 .ok()
8452 .and_then(|f| match f.data_type() {
8453 DataType::Union(uf, _) => uf
8454 .iter()
8455 .find(|(_, child)| child.name() == "uuid")
8456 .and_then(|(_, child)| {
8457 let md = child.metadata();
8458 let has_ext = md.get(UUID_EXT_KEY).is_some();
8459 let is_uuid_logical = md
8460 .get(UUID_LOGICAL_KEY)
8461 .map(|v| v.trim_matches('"') == "uuid")
8462 .unwrap_or(false);
8463 if has_ext || is_uuid_logical {
8464 Some(md.clone())
8465 } else {
8466 None
8467 }
8468 }),
8469 _ => None,
8470 });
8471
8472 let add_uuid_ext_top = |f: Field| -> Field {
8473 if let Some(md) = &uuid_md_top {
8474 f.with_metadata(md.clone())
8475 } else {
8476 f
8477 }
8478 };
8479 let add_uuid_ext_union = |f: Field| -> Field {
8480 if let Some(md) = &uuid_md_union {
8481 f.with_metadata(md.clone())
8482 } else {
8483 f
8484 }
8485 };
8486
8487 #[inline]
8488 fn uuid16_from_str(s: &str) -> [u8; 16] {
8489 let mut out = [0u8; 16];
8490 let mut idx = 0usize;
8491 let mut hi: Option<u8> = None;
8492 for ch in s.chars() {
8493 if ch == '-' {
8494 continue;
8495 }
8496 let v = ch.to_digit(16).expect("invalid hex digit in UUID") as u8;
8497 if let Some(h) = hi {
8498 out[idx] = (h << 4) | v;
8499 idx += 1;
8500 hi = None;
8501 } else {
8502 hi = Some(v);
8503 }
8504 }
8505 assert_eq!(idx, 16, "UUID must decode to 16 bytes");
8506 out
8507 }
8508
8509 fn mk_dense_union(
8510 fields: &UnionFields,
8511 type_ids: Vec<i8>,
8512 offsets: Vec<i32>,
8513 provide: impl Fn(&Field) -> Option<ArrayRef>,
8514 ) -> ArrayRef {
8515 fn empty_child_for(dt: &DataType) -> Arc<dyn Array> {
8516 match dt {
8517 DataType::Null => Arc::new(NullArray::new(0)),
8518 DataType::Boolean => Arc::new(BooleanArray::from(Vec::<bool>::new())),
8519 DataType::Int32 => Arc::new(Int32Array::from(Vec::<i32>::new())),
8520 DataType::Int64 => Arc::new(Int64Array::from(Vec::<i64>::new())),
8521 DataType::Float32 => Arc::new(Float32Array::from(Vec::<f32>::new())),
8522 DataType::Float64 => Arc::new(Float64Array::from(Vec::<f64>::new())),
8523 DataType::Binary => Arc::new(BinaryArray::from(Vec::<&[u8]>::new())),
8524 DataType::Utf8 => Arc::new(StringArray::from(Vec::<&str>::new())),
8525 DataType::Date32 => Arc::new(Date32Array::from(Vec::<i32>::new())),
8526 DataType::Time32(arrow_schema::TimeUnit::Millisecond) => {
8527 Arc::new(Time32MillisecondArray::from(Vec::<i32>::new()))
8528 }
8529 DataType::Time64(arrow_schema::TimeUnit::Microsecond) => {
8530 Arc::new(Time64MicrosecondArray::from(Vec::<i64>::new()))
8531 }
8532 DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, tz) => {
8533 let a = TimestampMillisecondArray::from(Vec::<i64>::new());
8534 Arc::new(if let Some(tz) = tz {
8535 a.with_timezone(tz.clone())
8536 } else {
8537 a
8538 })
8539 }
8540 DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, tz) => {
8541 let a = TimestampMicrosecondArray::from(Vec::<i64>::new());
8542 Arc::new(if let Some(tz) = tz {
8543 a.with_timezone(tz.clone())
8544 } else {
8545 a
8546 })
8547 }
8548 DataType::Interval(IntervalUnit::MonthDayNano) => Arc::new(
8549 IntervalMonthDayNanoArray::from(Vec::<IntervalMonthDayNano>::new()),
8550 ),
8551 DataType::FixedSizeBinary(sz) => Arc::new(
8552 FixedSizeBinaryArray::try_from_sparse_iter_with_size(
8553 std::iter::empty::<Option<Vec<u8>>>(),
8554 *sz,
8555 )
8556 .unwrap(),
8557 ),
8558 DataType::Dictionary(_, _) => {
8559 let keys = Int32Array::from(Vec::<i32>::new());
8560 let values = Arc::new(StringArray::from(Vec::<&str>::new()));
8561 Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap())
8562 }
8563 DataType::Struct(fields) => {
8564 let children: Vec<ArrayRef> = fields
8565 .iter()
8566 .map(|f| empty_child_for(f.data_type()) as ArrayRef)
8567 .collect();
8568 Arc::new(StructArray::new(fields.clone(), children, None))
8569 }
8570 DataType::List(field) => {
8571 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0]));
8572 Arc::new(
8573 ListArray::try_new(
8574 field.clone(),
8575 offsets,
8576 empty_child_for(field.data_type()),
8577 None,
8578 )
8579 .unwrap(),
8580 )
8581 }
8582 DataType::Map(entry_field, is_sorted) => {
8583 let (key_field, val_field) = match entry_field.data_type() {
8584 DataType::Struct(fs) => (fs[0].clone(), fs[1].clone()),
8585 other => panic!("unexpected map entries type: {other:?}"),
8586 };
8587 let keys = StringArray::from(Vec::<&str>::new());
8588 let vals: ArrayRef = match val_field.data_type() {
8589 DataType::Null => Arc::new(NullArray::new(0)) as ArrayRef,
8590 DataType::Boolean => {
8591 Arc::new(BooleanArray::from(Vec::<bool>::new())) as ArrayRef
8592 }
8593 DataType::Int32 => {
8594 Arc::new(Int32Array::from(Vec::<i32>::new())) as ArrayRef
8595 }
8596 DataType::Int64 => {
8597 Arc::new(Int64Array::from(Vec::<i64>::new())) as ArrayRef
8598 }
8599 DataType::Float32 => {
8600 Arc::new(Float32Array::from(Vec::<f32>::new())) as ArrayRef
8601 }
8602 DataType::Float64 => {
8603 Arc::new(Float64Array::from(Vec::<f64>::new())) as ArrayRef
8604 }
8605 DataType::Utf8 => {
8606 Arc::new(StringArray::from(Vec::<&str>::new())) as ArrayRef
8607 }
8608 DataType::Binary => {
8609 Arc::new(BinaryArray::from(Vec::<&[u8]>::new())) as ArrayRef
8610 }
8611 DataType::Union(uf, _) => {
8612 let children: Vec<ArrayRef> = uf
8613 .iter()
8614 .map(|(_, f)| empty_child_for(f.data_type()))
8615 .collect();
8616 Arc::new(
8617 UnionArray::try_new(
8618 uf.clone(),
8619 ScalarBuffer::<i8>::from(Vec::<i8>::new()),
8620 Some(ScalarBuffer::<i32>::from(Vec::<i32>::new())),
8621 children,
8622 )
8623 .unwrap(),
8624 ) as ArrayRef
8625 }
8626 other => panic!("unsupported map value type: {other:?}"),
8627 };
8628 let entries = StructArray::new(
8629 Fields::from(vec![
8630 key_field.as_ref().clone(),
8631 val_field.as_ref().clone(),
8632 ]),
8633 vec![Arc::new(keys) as ArrayRef, vals],
8634 None,
8635 );
8636 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0]));
8637 Arc::new(MapArray::new(
8638 entry_field.clone(),
8639 offsets,
8640 entries,
8641 None,
8642 *is_sorted,
8643 ))
8644 }
8645 other => panic!("empty_child_for: unhandled type {other:?}"),
8646 }
8647 }
8648 let children: Vec<ArrayRef> = fields
8649 .iter()
8650 .map(|(_, f)| provide(f).unwrap_or_else(|| empty_child_for(f.data_type())))
8651 .collect();
8652 Arc::new(
8653 UnionArray::try_new(
8654 fields.clone(),
8655 ScalarBuffer::<i8>::from(type_ids),
8656 Some(ScalarBuffer::<i32>::from(offsets)),
8657 children,
8658 )
8659 .unwrap(),
8660 ) as ArrayRef
8661 }
8662 let date_a: i32 = 19_000; let time_ms_a: i32 = 12 * 3_600_000 + 34 * 60_000 + 56_000 + 789;
8664 let time_us_eod: i64 = 86_400_000_000 - 1;
8665 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;
8667 let dur_small = IntervalMonthDayNanoType::make_value(1, 2, 3_000_000_000);
8668 let dur_zero = IntervalMonthDayNanoType::make_value(0, 0, 0);
8669 let dur_large =
8670 IntervalMonthDayNanoType::make_value(12, 31, ((86_400_000 - 1) as i64) * 1_000_000);
8671 let dur_2years = IntervalMonthDayNanoType::make_value(24, 0, 0);
8672 let uuid1 = uuid16_from_str("fe7bc30b-4ce8-4c5e-b67c-2234a2d38e66");
8673 let uuid2 = uuid16_from_str("0826cc06-d2e3-4599-b4ad-af5fa6905cdb");
8674 let item_name = Field::LIST_FIELD_DEFAULT_NAME;
8675 let uf_tri = UnionFields::try_new(
8676 vec![0, 1, 2],
8677 vec![
8678 Field::new("int", DataType::Int32, false),
8679 Field::new("string", DataType::Utf8, false),
8680 Field::new("boolean", DataType::Boolean, false),
8681 ],
8682 )
8683 .unwrap();
8684 let uf_arr_items = UnionFields::try_new(
8685 vec![0, 1, 2],
8686 vec![
8687 Field::new("null", DataType::Null, false),
8688 Field::new("string", DataType::Utf8, false),
8689 Field::new("long", DataType::Int64, false),
8690 ],
8691 )
8692 .unwrap();
8693 let arr_items_field = Arc::new(Field::new(
8694 item_name,
8695 DataType::Union(uf_arr_items.clone(), UnionMode::Dense),
8696 true,
8697 ));
8698 let uf_map_vals = UnionFields::try_new(
8699 vec![0, 1, 2],
8700 vec![
8701 Field::new("string", DataType::Utf8, false),
8702 Field::new("double", DataType::Float64, false),
8703 Field::new("null", DataType::Null, false),
8704 ],
8705 )
8706 .unwrap();
8707 let map_entries_field = Arc::new(Field::new(
8708 "entries",
8709 DataType::Struct(Fields::from(vec![
8710 Field::new("key", DataType::Utf8, false),
8711 Field::new(
8712 "value",
8713 DataType::Union(uf_map_vals.clone(), UnionMode::Dense),
8714 true,
8715 ),
8716 ])),
8717 false,
8718 ));
8719 let mut enum_md_color = {
8721 let mut m = HashMap::<String, String>::new();
8722 m.insert(
8723 crate::schema::AVRO_ENUM_SYMBOLS_METADATA_KEY.to_string(),
8724 serde_json::to_string(&vec!["RED", "GREEN", "BLUE"]).unwrap(),
8725 );
8726 m
8727 };
8728 enum_md_color.insert(AVRO_NAME_METADATA_KEY.to_string(), "Color".to_string());
8729 enum_md_color.insert(
8730 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8731 "org.apache.arrow.avrotests.v1.types".to_string(),
8732 );
8733 let union_rec_a_fields = Fields::from(vec![
8734 Field::new("a", DataType::Int32, false),
8735 Field::new("b", DataType::Utf8, false),
8736 ]);
8737 let union_rec_b_fields = Fields::from(vec![
8738 Field::new("x", DataType::Int64, false),
8739 Field::new("y", DataType::Binary, false),
8740 ]);
8741 let union_map_entries = Arc::new(Field::new(
8742 "entries",
8743 DataType::Struct(Fields::from(vec![
8744 Field::new("key", DataType::Utf8, false),
8745 Field::new("value", DataType::Utf8, false),
8746 ])),
8747 false,
8748 ));
8749 let person_md = {
8750 let mut m = HashMap::<String, String>::new();
8751 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "Person".to_string());
8752 m.insert(
8753 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8754 "com.example".to_string(),
8755 );
8756 m
8757 };
8758 let maybe_auth_md = {
8759 let mut m = HashMap::<String, String>::new();
8760 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "MaybeAuth".to_string());
8761 m.insert(
8762 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8763 "org.apache.arrow.avrotests.v1.types".to_string(),
8764 );
8765 m
8766 };
8767 let address_md = {
8768 let mut m = HashMap::<String, String>::new();
8769 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "Address".to_string());
8770 m.insert(
8771 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8772 "org.apache.arrow.avrotests.v1.types".to_string(),
8773 );
8774 m
8775 };
8776 let rec_a_md = {
8777 let mut m = HashMap::<String, String>::new();
8778 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "RecA".to_string());
8779 m.insert(
8780 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8781 "org.apache.arrow.avrotests.v1.types".to_string(),
8782 );
8783 m
8784 };
8785 let rec_b_md = {
8786 let mut m = HashMap::<String, String>::new();
8787 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "RecB".to_string());
8788 m.insert(
8789 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8790 "org.apache.arrow.avrotests.v1.types".to_string(),
8791 );
8792 m
8793 };
8794 let uf_union_big = UnionFields::try_new(
8795 vec![0, 1, 2, 3, 4],
8796 vec![
8797 Field::new(
8798 "map",
8799 DataType::Map(union_map_entries.clone(), false),
8800 false,
8801 ),
8802 Field::new(
8803 "array",
8804 DataType::List(Arc::new(Field::new(item_name, DataType::Int64, false))),
8805 false,
8806 ),
8807 Field::new(
8808 "org.apache.arrow.avrotests.v1.types.RecB",
8809 DataType::Struct(union_rec_b_fields.clone()),
8810 false,
8811 )
8812 .with_metadata(rec_b_md.clone()),
8813 Field::new(
8814 "org.apache.arrow.avrotests.v1.types.RecA",
8815 DataType::Struct(union_rec_a_fields.clone()),
8816 false,
8817 )
8818 .with_metadata(rec_a_md.clone()),
8819 Field::new(
8820 "org.apache.arrow.avrotests.v1.types.Color",
8821 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
8822 false,
8823 )
8824 .with_metadata(enum_md_color.clone()),
8825 ],
8826 )
8827 .unwrap();
8828 let fx4_md = {
8829 let mut m = HashMap::<String, String>::new();
8830 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "Fx4".to_string());
8831 m.insert(
8832 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8833 "org.apache.arrow.avrotests.v1".to_string(),
8834 );
8835 m
8836 };
8837 let uf_date_fixed4 = UnionFields::try_new(
8838 vec![0, 1],
8839 vec![
8840 Field::new(
8841 "org.apache.arrow.avrotests.v1.Fx4",
8842 DataType::FixedSizeBinary(4),
8843 false,
8844 )
8845 .with_metadata(fx4_md.clone()),
8846 Field::new("date", DataType::Date32, false),
8847 ],
8848 )
8849 .unwrap();
8850 let dur12u_md = {
8851 let mut m = HashMap::<String, String>::new();
8852 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "Dur12U".to_string());
8853 m.insert(
8854 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8855 "org.apache.arrow.avrotests.v1".to_string(),
8856 );
8857 m
8858 };
8859 let uf_dur_or_str = UnionFields::try_new(
8860 vec![0, 1],
8861 vec![
8862 Field::new("string", DataType::Utf8, false),
8863 Field::new(
8864 "org.apache.arrow.avrotests.v1.Dur12U",
8865 DataType::Interval(arrow_schema::IntervalUnit::MonthDayNano),
8866 false,
8867 )
8868 .with_metadata(dur12u_md.clone()),
8869 ],
8870 )
8871 .unwrap();
8872 let fx10_md = {
8873 let mut m = HashMap::<String, String>::new();
8874 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "Fx10".to_string());
8875 m.insert(
8876 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8877 "org.apache.arrow.avrotests.v1".to_string(),
8878 );
8879 m
8880 };
8881 let uf_uuid_or_fx10 = UnionFields::try_new(
8882 vec![0, 1],
8883 vec![
8884 Field::new(
8885 "org.apache.arrow.avrotests.v1.Fx10",
8886 DataType::FixedSizeBinary(10),
8887 false,
8888 )
8889 .with_metadata(fx10_md.clone()),
8890 add_uuid_ext_union(Field::new("uuid", DataType::FixedSizeBinary(16), false)),
8891 ],
8892 )
8893 .unwrap();
8894 let uf_kv_val = UnionFields::try_new(
8895 vec![0, 1, 2],
8896 vec![
8897 Field::new("null", DataType::Null, false),
8898 Field::new("int", DataType::Int32, false),
8899 Field::new("long", DataType::Int64, false),
8900 ],
8901 )
8902 .unwrap();
8903 let kv_fields = Fields::from(vec![
8904 Field::new("key", DataType::Utf8, false),
8905 Field::new(
8906 "val",
8907 DataType::Union(uf_kv_val.clone(), UnionMode::Dense),
8908 true,
8909 ),
8910 ]);
8911 let kv_md = {
8912 let mut m = HashMap::<String, String>::new();
8913 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "KV".to_string());
8914 m.insert(
8915 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8916 "org.apache.arrow.avrotests.v1.types".to_string(),
8917 );
8918 m
8919 };
8920 let kv_item_field = Arc::new(
8921 Field::new(item_name, DataType::Struct(kv_fields.clone()), false).with_metadata(kv_md),
8922 );
8923 let map_int_entries = Arc::new(Field::new(
8924 "entries",
8925 DataType::Struct(Fields::from(vec![
8926 Field::new("key", DataType::Utf8, false),
8927 Field::new("value", DataType::Int32, false),
8928 ])),
8929 false,
8930 ));
8931 let uf_map_or_array = UnionFields::try_new(
8932 vec![0, 1],
8933 vec![
8934 Field::new(
8935 "array",
8936 DataType::List(Arc::new(Field::new(item_name, DataType::Int32, false))),
8937 false,
8938 ),
8939 Field::new("map", DataType::Map(map_int_entries.clone(), false), false),
8940 ],
8941 )
8942 .unwrap();
8943 let mut enum_md_status = {
8944 let mut m = HashMap::<String, String>::new();
8945 m.insert(
8946 crate::schema::AVRO_ENUM_SYMBOLS_METADATA_KEY.to_string(),
8947 serde_json::to_string(&vec!["UNKNOWN", "NEW", "PROCESSING", "DONE"]).unwrap(),
8948 );
8949 m
8950 };
8951 enum_md_status.insert(AVRO_NAME_METADATA_KEY.to_string(), "Status".to_string());
8952 enum_md_status.insert(
8953 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8954 "org.apache.arrow.avrotests.v1.types".to_string(),
8955 );
8956 let mut dec20_md = HashMap::<String, String>::new();
8957 dec20_md.insert("precision".to_string(), "20".to_string());
8958 dec20_md.insert("scale".to_string(), "4".to_string());
8959 dec20_md.insert(AVRO_NAME_METADATA_KEY.to_string(), "DecFix20".to_string());
8960 dec20_md.insert(
8961 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8962 "org.apache.arrow.avrotests.v1.types".to_string(),
8963 );
8964 let mut dec10_md = HashMap::<String, String>::new();
8965 dec10_md.insert("precision".to_string(), "10".to_string());
8966 dec10_md.insert("scale".to_string(), "2".to_string());
8967 let fx16_top_md = {
8968 let mut m = HashMap::<String, String>::new();
8969 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "Fx16".to_string());
8970 m.insert(
8971 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8972 "org.apache.arrow.avrotests.v1.types".to_string(),
8973 );
8974 m
8975 };
8976 let dur12_top_md = {
8977 let mut m = HashMap::<String, String>::new();
8978 m.insert(AVRO_NAME_METADATA_KEY.to_string(), "Dur12".to_string());
8979 m.insert(
8980 AVRO_NAMESPACE_METADATA_KEY.to_string(),
8981 "org.apache.arrow.avrotests.v1.types".to_string(),
8982 );
8983 m
8984 };
8985 #[cfg(feature = "small_decimals")]
8986 let dec20_dt = DataType::Decimal128(20, 4);
8987 #[cfg(not(feature = "small_decimals"))]
8988 let dec20_dt = DataType::Decimal128(20, 4);
8989 #[cfg(feature = "small_decimals")]
8990 let dec10_dt = DataType::Decimal64(10, 2);
8991 #[cfg(not(feature = "small_decimals"))]
8992 let dec10_dt = DataType::Decimal128(10, 2);
8993 let fields: Vec<FieldRef> = vec![
8994 Arc::new(
8995 Field::new(
8996 "person",
8997 DataType::Struct(Fields::from(vec![
8998 Field::new("name", DataType::Utf8, false),
8999 Field::new("age", DataType::Int32, false),
9000 ])),
9001 false,
9002 )
9003 .with_metadata(person_md),
9004 ),
9005 Arc::new(Field::new("old_count", DataType::Int32, false)),
9006 Arc::new(Field::new(
9007 "union_map_or_array_int",
9008 DataType::Union(uf_map_or_array.clone(), UnionMode::Dense),
9009 false,
9010 )),
9011 Arc::new(Field::new(
9012 "array_records_with_union",
9013 DataType::List(kv_item_field.clone()),
9014 false,
9015 )),
9016 Arc::new(Field::new(
9017 "union_uuid_or_fixed10",
9018 DataType::Union(uf_uuid_or_fx10.clone(), UnionMode::Dense),
9019 false,
9020 )),
9021 Arc::new(Field::new(
9022 "union_interval_or_string",
9023 DataType::Union(uf_dur_or_str.clone(), UnionMode::Dense),
9024 false,
9025 )),
9026 Arc::new(Field::new(
9027 "union_date_or_fixed4",
9028 DataType::Union(uf_date_fixed4.clone(), UnionMode::Dense),
9029 false,
9030 )),
9031 Arc::new(Field::new(
9032 "union_enum_record_array_map",
9033 DataType::Union(uf_union_big.clone(), UnionMode::Dense),
9034 false,
9035 )),
9036 Arc::new(
9037 Field::new(
9038 "maybe_auth",
9039 DataType::Struct(Fields::from(vec![
9040 Field::new("user", DataType::Utf8, false),
9041 Field::new("token", DataType::Binary, true), ])),
9043 false,
9044 )
9045 .with_metadata(maybe_auth_md),
9046 ),
9047 Arc::new(
9048 Field::new(
9049 "address",
9050 DataType::Struct(Fields::from(vec![
9051 Field::new("street_name", DataType::Utf8, false),
9052 Field::new("zip", DataType::Int32, false),
9053 Field::new("country", DataType::Utf8, false),
9054 ])),
9055 false,
9056 )
9057 .with_metadata(address_md),
9058 ),
9059 Arc::new(Field::new(
9060 "map_union",
9061 DataType::Map(map_entries_field.clone(), false),
9062 false,
9063 )),
9064 Arc::new(Field::new(
9065 "arr_union",
9066 DataType::List(arr_items_field.clone()),
9067 false,
9068 )),
9069 Arc::new(
9070 Field::new(
9071 "status",
9072 DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
9073 false,
9074 )
9075 .with_metadata(enum_md_status.clone()),
9076 ),
9077 Arc::new(
9078 Field::new(
9079 "interval_mdn",
9080 DataType::Interval(IntervalUnit::MonthDayNano),
9081 false,
9082 )
9083 .with_metadata(dur12_top_md.clone()),
9084 ),
9085 Arc::new(Field::new(
9086 "ts_micros_local",
9087 DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, None),
9088 false,
9089 )),
9090 Arc::new(Field::new(
9091 "ts_millis_local",
9092 DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, None),
9093 false,
9094 )),
9095 Arc::new(Field::new(
9096 "ts_micros_utc",
9097 DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, Some("+00:00".into())),
9098 false,
9099 )),
9100 Arc::new(Field::new(
9101 "ts_millis_utc",
9102 DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, Some("+00:00".into())),
9103 false,
9104 )),
9105 Arc::new(Field::new(
9106 "t_micros",
9107 DataType::Time64(arrow_schema::TimeUnit::Microsecond),
9108 false,
9109 )),
9110 Arc::new(Field::new(
9111 "t_millis",
9112 DataType::Time32(arrow_schema::TimeUnit::Millisecond),
9113 false,
9114 )),
9115 Arc::new(Field::new("d_date", DataType::Date32, false)),
9116 Arc::new(add_uuid_ext_top(Field::new(
9117 "uuid_str",
9118 DataType::FixedSizeBinary(16),
9119 false,
9120 ))),
9121 Arc::new(Field::new("dec_fix_s20_4", dec20_dt, false).with_metadata(dec20_md.clone())),
9122 Arc::new(
9123 Field::new("dec_bytes_s10_2", dec10_dt, false).with_metadata(dec10_md.clone()),
9124 ),
9125 Arc::new(
9126 Field::new("fx16_plain", DataType::FixedSizeBinary(16), false)
9127 .with_metadata(fx16_top_md.clone()),
9128 ),
9129 Arc::new(Field::new("raw_bytes", DataType::Binary, false)),
9130 Arc::new(Field::new("str_utf8", DataType::Utf8, false)),
9131 Arc::new(Field::new(
9132 "tri_union_prim",
9133 DataType::Union(uf_tri.clone(), UnionMode::Dense),
9134 false,
9135 )),
9136 Arc::new(Field::new("opt_str_nullsecond", DataType::Utf8, true)),
9137 Arc::new(Field::new("opt_i32_nullfirst", DataType::Int32, true)),
9138 Arc::new(Field::new("count_i64", DataType::Int64, false)),
9139 Arc::new(Field::new("count_i32", DataType::Int64, false)),
9140 Arc::new(Field::new("ratio_f64", DataType::Float64, false)),
9141 Arc::new(Field::new("ratio_f32", DataType::Float64, false)),
9142 Arc::new(Field::new("flag", DataType::Boolean, false)),
9143 Arc::new(Field::new("identifier", DataType::Int64, false)),
9144 ];
9145 let expected_schema = Arc::new(arrow_schema::Schema::new(Fields::from(fields)));
9146 let mut cols: Vec<ArrayRef> = vec![
9147 Arc::new(StructArray::new(
9148 match expected_schema
9149 .field_with_name("person")
9150 .unwrap()
9151 .data_type()
9152 {
9153 DataType::Struct(fs) => fs.clone(),
9154 _ => unreachable!(),
9155 },
9156 vec![
9157 Arc::new(StringArray::from(vec!["Alice", "Bob", "Carol", "Dave"])) as ArrayRef,
9158 Arc::new(Int32Array::from(vec![30, 0, 25, 41])) as ArrayRef,
9159 ],
9160 None,
9161 )) as ArrayRef,
9162 Arc::new(Int32Array::from(vec![100, 42, 7, 42])) as ArrayRef,
9163 ];
9164 {
9165 let map_child: ArrayRef = {
9166 let keys = StringArray::from(vec!["x", "y", "only"]);
9167 let vals = Int32Array::from(vec![1, 2, 10]);
9168 let entries = StructArray::new(
9169 Fields::from(vec![
9170 Field::new("key", DataType::Utf8, false),
9171 Field::new("value", DataType::Int32, false),
9172 ]),
9173 vec![Arc::new(keys) as ArrayRef, Arc::new(vals) as ArrayRef],
9174 None,
9175 );
9176 let moff = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2, 3]));
9177 Arc::new(MapArray::new(
9178 map_int_entries.clone(),
9179 moff,
9180 entries,
9181 None,
9182 false,
9183 )) as ArrayRef
9184 };
9185 let list_child: ArrayRef = {
9186 let values = Int32Array::from(vec![1, 2, 3, 0]);
9187 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 4]));
9188 Arc::new(
9189 ListArray::try_new(
9190 Arc::new(Field::new(item_name, DataType::Int32, false)),
9191 offsets,
9192 Arc::new(values),
9193 None,
9194 )
9195 .unwrap(),
9196 ) as ArrayRef
9197 };
9198 let tids = vec![1, 0, 1, 0];
9199 let offs = vec![0, 0, 1, 1];
9200 let arr = mk_dense_union(&uf_map_or_array, tids, offs, |f| match f.name().as_str() {
9201 "array" => Some(list_child.clone()),
9202 "map" => Some(map_child.clone()),
9203 _ => None,
9204 });
9205 cols.push(arr);
9206 }
9207 {
9208 let keys = Arc::new(StringArray::from(vec!["k1", "k2", "k", "k3", "x"])) as ArrayRef;
9209 let type_ids = vec![1, 0, 2, 0, 1];
9210 let offsets = vec![0, 0, 0, 1, 1];
9211 let vals = mk_dense_union(&uf_kv_val, type_ids, offsets, |f| match f.data_type() {
9212 DataType::Int32 => Some(Arc::new(Int32Array::from(vec![5, -5])) as ArrayRef),
9213 DataType::Int64 => Some(Arc::new(Int64Array::from(vec![99i64])) as ArrayRef),
9214 DataType::Null => Some(Arc::new(NullArray::new(2)) as ArrayRef),
9215 _ => None,
9216 });
9217 let values_struct =
9218 Arc::new(StructArray::new(kv_fields.clone(), vec![keys, vals], None));
9219 let list_offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 2, 3, 4, 5]));
9220 let arr = Arc::new(
9221 ListArray::try_new(kv_item_field.clone(), list_offsets, values_struct, None)
9222 .unwrap(),
9223 ) as ArrayRef;
9224 cols.push(arr);
9225 }
9226 {
9227 let type_ids = vec![1, 0, 1, 0]; let offs = vec![0, 0, 1, 1];
9229 let arr = mk_dense_union(&uf_uuid_or_fx10, type_ids, offs, |f| match f.data_type() {
9230 DataType::FixedSizeBinary(16) => {
9231 let it = [Some(uuid1), Some(uuid2)].into_iter();
9232 Some(Arc::new(
9233 FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 16).unwrap(),
9234 ) as ArrayRef)
9235 }
9236 DataType::FixedSizeBinary(10) => {
9237 let fx10_a = [0xAAu8; 10];
9238 let fx10_b = [0x00u8, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77, 0x88, 0x99];
9239 let it = [Some(fx10_a), Some(fx10_b)].into_iter();
9240 Some(Arc::new(
9241 FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 10).unwrap(),
9242 ) as ArrayRef)
9243 }
9244 _ => None,
9245 });
9246 cols.push(arr);
9247 }
9248 {
9249 let type_ids = vec![1, 0, 1, 0]; let offs = vec![0, 0, 1, 1];
9251 let arr = mk_dense_union(&uf_dur_or_str, type_ids, offs, |f| match f.data_type() {
9252 DataType::Interval(arrow_schema::IntervalUnit::MonthDayNano) => Some(Arc::new(
9253 IntervalMonthDayNanoArray::from(vec![dur_small, dur_large]),
9254 )
9255 as ArrayRef),
9256 DataType::Utf8 => Some(Arc::new(StringArray::from(vec![
9257 "duration-as-text",
9258 "iso-8601-period-P1Y",
9259 ])) as ArrayRef),
9260 _ => None,
9261 });
9262 cols.push(arr);
9263 }
9264 {
9265 let type_ids = vec![1, 0, 1, 0]; let offs = vec![0, 0, 1, 1];
9267 let arr = mk_dense_union(&uf_date_fixed4, type_ids, offs, |f| match f.data_type() {
9268 DataType::Date32 => Some(Arc::new(Date32Array::from(vec![date_a, 0])) as ArrayRef),
9269 DataType::FixedSizeBinary(4) => {
9270 let it = [Some(*b"\x00\x11\x22\x33"), Some(*b"ABCD")].into_iter();
9271 Some(Arc::new(
9272 FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 4).unwrap(),
9273 ) as ArrayRef)
9274 }
9275 _ => None,
9276 });
9277 cols.push(arr);
9278 }
9279 {
9280 let tids = vec![4, 3, 1, 0]; let offs = vec![0, 0, 0, 0];
9282 let arr = mk_dense_union(&uf_union_big, tids, offs, |f| match f.data_type() {
9283 DataType::Dictionary(_, _) => {
9284 let keys = Int32Array::from(vec![0i32]);
9285 let values =
9286 Arc::new(StringArray::from(vec!["RED", "GREEN", "BLUE"])) as ArrayRef;
9287 Some(
9288 Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap())
9289 as ArrayRef,
9290 )
9291 }
9292 DataType::Struct(fs) if fs == &union_rec_a_fields => {
9293 let a = Int32Array::from(vec![7]);
9294 let b = StringArray::from(vec!["rec"]);
9295 Some(Arc::new(StructArray::new(
9296 fs.clone(),
9297 vec![Arc::new(a) as ArrayRef, Arc::new(b) as ArrayRef],
9298 None,
9299 )) as ArrayRef)
9300 }
9301 DataType::List(_) => {
9302 let values = Int64Array::from(vec![1i64, 2, 3]);
9303 let offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3]));
9304 Some(Arc::new(
9305 ListArray::try_new(
9306 Arc::new(Field::new(item_name, DataType::Int64, false)),
9307 offsets,
9308 Arc::new(values),
9309 None,
9310 )
9311 .unwrap(),
9312 ) as ArrayRef)
9313 }
9314 DataType::Map(_, _) => {
9315 let keys = StringArray::from(vec!["k"]);
9316 let vals = StringArray::from(vec!["v"]);
9317 let entries = StructArray::new(
9318 Fields::from(vec![
9319 Field::new("key", DataType::Utf8, false),
9320 Field::new("value", DataType::Utf8, false),
9321 ]),
9322 vec![Arc::new(keys) as ArrayRef, Arc::new(vals) as ArrayRef],
9323 None,
9324 );
9325 let moff = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 1]));
9326 Some(Arc::new(MapArray::new(
9327 union_map_entries.clone(),
9328 moff,
9329 entries,
9330 None,
9331 false,
9332 )) as ArrayRef)
9333 }
9334 _ => None,
9335 });
9336 cols.push(arr);
9337 }
9338 {
9339 let fs = match expected_schema
9340 .field_with_name("maybe_auth")
9341 .unwrap()
9342 .data_type()
9343 {
9344 DataType::Struct(fs) => fs.clone(),
9345 _ => unreachable!(),
9346 };
9347 let user =
9348 Arc::new(StringArray::from(vec!["alice", "bob", "carol", "dave"])) as ArrayRef;
9349 let token_values: Vec<Option<&[u8]>> = vec![
9350 None,
9351 Some(b"\x01\x02\x03".as_ref()),
9352 None,
9353 Some(b"".as_ref()),
9354 ];
9355 let token = Arc::new(BinaryArray::from(token_values)) as ArrayRef;
9356 cols.push(Arc::new(StructArray::new(fs, vec![user, token], None)) as ArrayRef);
9357 }
9358 {
9359 let fs = match expected_schema
9360 .field_with_name("address")
9361 .unwrap()
9362 .data_type()
9363 {
9364 DataType::Struct(fs) => fs.clone(),
9365 _ => unreachable!(),
9366 };
9367 let street = Arc::new(StringArray::from(vec![
9368 "100 Main",
9369 "",
9370 "42 Galaxy Way",
9371 "End Ave",
9372 ])) as ArrayRef;
9373 let zip = Arc::new(Int32Array::from(vec![12345, 0, 42424, 1])) as ArrayRef;
9374 let country = Arc::new(StringArray::from(vec!["US", "CA", "US", "GB"])) as ArrayRef;
9375 cols.push(Arc::new(StructArray::new(fs, vec![street, zip, country], None)) as ArrayRef);
9376 }
9377 {
9378 let keys = StringArray::from(vec!["a", "b", "c", "neg", "pi", "ok"]);
9379 let moff = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 3, 4, 4, 6]));
9380 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];
9384 let offsets = vec![0, 0, 0, 1, 2, 1];
9385 let pi_5dp = (std::f64::consts::PI * 100_000.0).trunc() / 100_000.0;
9386 let vals = mk_dense_union(&uf_map_vals, type_ids, offsets, |f| match f.data_type() {
9387 DataType::Float64 => {
9388 Some(Arc::new(Float64Array::from(vec![1.5f64, -0.5, pi_5dp])) as ArrayRef)
9389 }
9390 DataType::Utf8 => {
9391 Some(Arc::new(StringArray::from(vec!["yes", "true"])) as ArrayRef)
9392 }
9393 DataType::Null => Some(Arc::new(NullArray::new(1)) as ArrayRef),
9394 _ => None,
9395 });
9396 let entries = StructArray::new(
9397 Fields::from(vec![
9398 Field::new("key", DataType::Utf8, false),
9399 Field::new(
9400 "value",
9401 DataType::Union(uf_map_vals.clone(), UnionMode::Dense),
9402 true,
9403 ),
9404 ]),
9405 vec![Arc::new(keys) as ArrayRef, vals],
9406 None,
9407 );
9408 let map = Arc::new(MapArray::new(
9409 map_entries_field.clone(),
9410 moff,
9411 entries,
9412 None,
9413 false,
9414 )) as ArrayRef;
9415 cols.push(map);
9416 }
9417 {
9418 let type_ids = vec![
9419 2, 1, 0, 2, 0, 1, 2, 2, 1, 0,
9420 2, ];
9422 let offsets = vec![0, 0, 0, 1, 1, 1, 2, 3, 2, 2, 4];
9423 let values =
9424 mk_dense_union(&uf_arr_items, type_ids, offsets, |f| match f.data_type() {
9425 DataType::Int64 => {
9426 Some(Arc::new(Int64Array::from(vec![1i64, -3, 0, -1, 0])) as ArrayRef)
9427 }
9428 DataType::Utf8 => {
9429 Some(Arc::new(StringArray::from(vec!["x", "z", "end"])) as ArrayRef)
9430 }
9431 DataType::Null => Some(Arc::new(NullArray::new(3)) as ArrayRef),
9432 _ => None,
9433 });
9434 let list_offsets = OffsetBuffer::new(ScalarBuffer::<i32>::from(vec![0, 4, 7, 8, 11]));
9435 let arr = Arc::new(
9436 ListArray::try_new(arr_items_field.clone(), list_offsets, values, None).unwrap(),
9437 ) as ArrayRef;
9438 cols.push(arr);
9439 }
9440 {
9441 let keys = Int32Array::from(vec![1, 2, 3, 0]); let values = Arc::new(StringArray::from(vec![
9443 "UNKNOWN",
9444 "NEW",
9445 "PROCESSING",
9446 "DONE",
9447 ])) as ArrayRef;
9448 let dict = DictionaryArray::<Int32Type>::try_new(keys, values).unwrap();
9449 cols.push(Arc::new(dict) as ArrayRef);
9450 }
9451 cols.push(Arc::new(IntervalMonthDayNanoArray::from(vec![
9452 dur_small, dur_zero, dur_large, dur_2years,
9453 ])) as ArrayRef);
9454 cols.push(Arc::new(TimestampMicrosecondArray::from(vec![
9455 ts_us_2024_01_01 + 123_456,
9456 0,
9457 ts_us_2024_01_01 + 101_112,
9458 987_654_321,
9459 ])) as ArrayRef);
9460 cols.push(Arc::new(TimestampMillisecondArray::from(vec![
9461 ts_ms_2024_01_01 + 86_400_000,
9462 0,
9463 ts_ms_2024_01_01 + 789,
9464 123_456_789,
9465 ])) as ArrayRef);
9466 {
9467 let a = TimestampMicrosecondArray::from(vec![
9468 ts_us_2024_01_01,
9469 1,
9470 ts_us_2024_01_01 + 456,
9471 0,
9472 ])
9473 .with_timezone("+00:00");
9474 cols.push(Arc::new(a) as ArrayRef);
9475 }
9476 {
9477 let a = TimestampMillisecondArray::from(vec![
9478 ts_ms_2024_01_01,
9479 -1,
9480 ts_ms_2024_01_01 + 123,
9481 0,
9482 ])
9483 .with_timezone("+00:00");
9484 cols.push(Arc::new(a) as ArrayRef);
9485 }
9486 cols.push(Arc::new(Time64MicrosecondArray::from(vec![
9487 time_us_eod,
9488 0,
9489 1,
9490 1_000_000,
9491 ])) as ArrayRef);
9492 cols.push(Arc::new(Time32MillisecondArray::from(vec![
9493 time_ms_a,
9494 0,
9495 1,
9496 86_400_000 - 1,
9497 ])) as ArrayRef);
9498 cols.push(Arc::new(Date32Array::from(vec![date_a, 0, 1, 365])) as ArrayRef);
9499 {
9500 let it = [Some(uuid1), Some(uuid2), Some(uuid1), Some(uuid2)].into_iter();
9501 cols.push(Arc::new(
9502 FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 16).unwrap(),
9503 ) as ArrayRef);
9504 }
9505 {
9506 #[cfg(feature = "small_decimals")]
9507 let arr = Arc::new(
9508 Decimal128Array::from_iter_values([1_234_567_891_234i128, -420_000i128, 0, -1i128])
9509 .with_precision_and_scale(20, 4)
9510 .unwrap(),
9511 ) as ArrayRef;
9512 #[cfg(not(feature = "small_decimals"))]
9513 let arr = Arc::new(
9514 Decimal128Array::from_iter_values([1_234_567_891_234i128, -420_000i128, 0, -1i128])
9515 .with_precision_and_scale(20, 4)
9516 .unwrap(),
9517 ) as ArrayRef;
9518 cols.push(arr);
9519 }
9520 {
9521 #[cfg(feature = "small_decimals")]
9522 let arr = Arc::new(
9523 Decimal64Array::from_iter_values([123456i64, -1, 0, 9_999_999_999i64])
9524 .with_precision_and_scale(10, 2)
9525 .unwrap(),
9526 ) as ArrayRef;
9527 #[cfg(not(feature = "small_decimals"))]
9528 let arr = Arc::new(
9529 Decimal128Array::from_iter_values([123456i128, -1, 0, 9_999_999_999i128])
9530 .with_precision_and_scale(10, 2)
9531 .unwrap(),
9532 ) as ArrayRef;
9533 cols.push(arr);
9534 }
9535 {
9536 let it = [
9537 Some(*b"0123456789ABCDEF"),
9538 Some([0u8; 16]),
9539 Some(*b"ABCDEFGHIJKLMNOP"),
9540 Some([0xAA; 16]),
9541 ]
9542 .into_iter();
9543 cols.push(Arc::new(
9544 FixedSizeBinaryArray::try_from_sparse_iter_with_size(it, 16).unwrap(),
9545 ) as ArrayRef);
9546 }
9547 cols.push(Arc::new(BinaryArray::from(vec![
9548 b"\x00\x01".as_ref(),
9549 b"".as_ref(),
9550 b"\xFF\x00".as_ref(),
9551 b"\x10\x20\x30\x40".as_ref(),
9552 ])) as ArrayRef);
9553 cols.push(Arc::new(StringArray::from(vec!["hello", "", "world", "✓ unicode"])) as ArrayRef);
9554 {
9555 let tids = vec![0, 1, 2, 1];
9556 let offs = vec![0, 0, 0, 1];
9557 let arr = mk_dense_union(&uf_tri, tids, offs, |f| match f.data_type() {
9558 DataType::Int32 => Some(Arc::new(Int32Array::from(vec![0])) as ArrayRef),
9559 DataType::Utf8 => Some(Arc::new(StringArray::from(vec!["hi", ""])) as ArrayRef),
9560 DataType::Boolean => Some(Arc::new(BooleanArray::from(vec![true])) as ArrayRef),
9561 _ => None,
9562 });
9563 cols.push(arr);
9564 }
9565 cols.push(Arc::new(StringArray::from(vec![
9566 Some("alpha"),
9567 None,
9568 Some("s3"),
9569 Some(""),
9570 ])) as ArrayRef);
9571 cols.push(Arc::new(Int32Array::from(vec![None, Some(42), None, Some(0)])) as ArrayRef);
9572 cols.push(Arc::new(Int64Array::from(vec![
9573 7_000_000_000i64,
9574 -2,
9575 0,
9576 -9_876_543_210i64,
9577 ])) as ArrayRef);
9578 cols.push(Arc::new(Int64Array::from(vec![7i64, -1, 0, 123])) as ArrayRef);
9579 cols.push(Arc::new(Float64Array::from(vec![2.5f64, -1.0, 7.0, -2.25])) as ArrayRef);
9580 cols.push(Arc::new(Float64Array::from(vec![1.25f64, -0.0, 3.5, 9.75])) as ArrayRef);
9581 cols.push(Arc::new(BooleanArray::from(vec![true, false, true, false])) as ArrayRef);
9582 cols.push(Arc::new(Int64Array::from(vec![1, 2, 3, 4])) as ArrayRef);
9583 let expected = RecordBatch::try_new(expected_schema, cols).unwrap();
9584 assert_eq!(
9585 expected, batch,
9586 "entire RecordBatch mismatch (schema, all columns, all rows)"
9587 );
9588 }
9589
9590 fn make_type_ref_ocf() -> Vec<u8> {
9600 use apache_avro::{Schema as ApacheSchema, Writer as ApacheWriter, types::Value};
9601 let schema_json = r#"{
9602 "type": "record", "name": "Root",
9603 "fields": [
9604 {"name": "ts", "type": {"type": "record", "name": "Timestamp", "fields": [
9605 {"name": "seconds", "type": "long"},
9606 {"name": "nanos", "type": "int"}
9607 ]}},
9608 {"name": "extra", "type": {"type": "record", "name": "Event", "fields": [
9609 {"name": "time", "type": "Timestamp"}
9610 ]}}
9611 ]
9612 }"#;
9613 let schema = ApacheSchema::parse_str(schema_json).expect("valid schema");
9614 let mut out = Vec::new();
9615 {
9616 let mut writer = ApacheWriter::new(&schema, &mut out);
9617 let ts_val = |s: i64, n: i32| {
9618 Value::Record(vec![
9619 ("seconds".into(), Value::Long(s)),
9620 ("nanos".into(), Value::Int(n)),
9621 ])
9622 };
9623 for (ts_s, ts_n, ex_s, ex_n) in [(1000i64, 100i32, -1i64, -1i32), (2000, 200, -2, -2)] {
9625 let row = Value::Record(vec![
9626 ("ts".into(), ts_val(ts_s, ts_n)),
9627 (
9628 "extra".into(),
9629 Value::Record(vec![("time".into(), ts_val(ex_s, ex_n))]),
9630 ),
9631 ]);
9632 writer.append_value_ref(&row).expect("append row");
9633 }
9634 writer.flush().expect("flush");
9635 }
9636 out
9637 }
9638
9639 #[test]
9648 fn test_nullable_reader_schema_vs_plain_writer_nested_struct() {
9649 let bytes = make_type_ref_ocf();
9650 let reader_schema = AvroSchema::new(
9651 r#"{"type":"record","name":"Root","fields":[
9652 {"name":"ts","type":["null",{"type":"record","name":"Timestamp","fields":[
9653 {"name":"seconds","type":["null","long"]},
9654 {"name":"nanos", "type":["null","int"]}
9655 ]}]}
9656 ]}"#
9657 .to_string(),
9658 );
9659 let mut reader = ReaderBuilder::new()
9660 .with_reader_schema(reader_schema)
9661 .build(Cursor::new(bytes))
9662 .expect("reader should build");
9663 let batch = reader
9664 .next()
9665 .expect("should have a batch")
9666 .expect("reading should succeed");
9667 assert_eq!(batch.num_rows(), 2);
9668 let ts = batch
9669 .column(0)
9670 .as_any()
9671 .downcast_ref::<StructArray>()
9672 .unwrap();
9673 let seconds = ts
9674 .column_by_name("seconds")
9675 .unwrap()
9676 .as_any()
9677 .downcast_ref::<Int64Array>()
9678 .unwrap();
9679 assert_eq!(seconds.value(0), 1000);
9680 assert_eq!(seconds.value(1), 2000);
9681 }
9682
9683 #[test]
9691 fn test_skipper_consumes_writer_only_struct_fields() {
9692 let bytes = make_type_ref_ocf();
9693 let reader_schema = AvroSchema::new(
9694 r#"{"type":"record","name":"Root","fields":[
9695 {"name":"ts","type":{"type":"record","name":"Timestamp","fields":[
9696 {"name":"seconds","type":"long"}
9697 ]}}
9698 ]}"#
9699 .to_string(),
9700 );
9701 let mut reader = ReaderBuilder::new()
9702 .with_reader_schema(reader_schema)
9703 .build(Cursor::new(bytes))
9704 .expect("reader should build");
9705 let batch = reader
9706 .next()
9707 .expect("should have a batch")
9708 .expect("Skipper must consume both seconds and nanos for extra.time");
9709 assert_eq!(batch.num_rows(), 2);
9710 let ts = batch
9711 .column(0)
9712 .as_any()
9713 .downcast_ref::<StructArray>()
9714 .unwrap();
9715 let seconds = ts
9716 .column_by_name("seconds")
9717 .unwrap()
9718 .as_any()
9719 .downcast_ref::<Int64Array>()
9720 .unwrap();
9721 assert_eq!(seconds.value(0), 1000);
9722 assert_eq!(seconds.value(1), 2000);
9723 }
9724
9725 #[test]
9734 fn test_skip_array_of_structs_uses_writer_schema_not_resolved() {
9735 use apache_avro::{Schema as ApacheSchema, Writer as ApacheWriter, types::Value};
9736 let schema_json = r#"{
9737 "type": "record", "name": "Root",
9738 "fields": [
9739 {"name": "ts", "type": {"type": "record", "name": "Timestamp", "fields": [
9740 {"name": "seconds", "type": "long"},
9741 {"name": "nanos", "type": "int"}
9742 ]}},
9743 {"name": "events", "type": {"type": "array", "items": {
9744 "type": "record", "name": "Event", "fields": [
9745 {"name": "time", "type": "Timestamp"}
9746 ]
9747 }}}
9748 ]
9749 }"#;
9750 let schema = ApacheSchema::parse_str(schema_json).expect("valid schema");
9751 let mut bytes = Vec::new();
9752 {
9753 let mut writer = ApacheWriter::new(&schema, &mut bytes);
9754 let ts_val = |s: i64, n: i32| {
9756 Value::Record(vec![
9757 ("seconds".into(), Value::Long(s)),
9758 ("nanos".into(), Value::Int(n)),
9759 ])
9760 };
9761 let row = Value::Record(vec![
9762 ("ts".into(), ts_val(100, 5)),
9763 (
9764 "events".into(),
9765 Value::Array(vec![Value::Record(vec![("time".into(), ts_val(200, 1))])]),
9766 ),
9767 ]);
9768 writer.append_value_ref(&row).expect("append row");
9769 writer.flush().expect("flush");
9770 }
9771
9772 let reader_schema = AvroSchema::new(
9774 r#"{"type":"record","name":"Root","fields":[
9775 {"name":"ts","type":["null",{"type":"record","name":"Timestamp","fields":[
9776 {"name":"seconds","type":["null","long"]},
9777 {"name":"nanos", "type":["null","int"]}
9778 ]}]}
9779 ]}"#
9780 .to_string(),
9781 );
9782 let mut reader = ReaderBuilder::new()
9783 .with_reader_schema(reader_schema)
9784 .build(Cursor::new(bytes))
9785 .expect("reader should build");
9786 let batch = reader
9787 .next()
9788 .expect("should have a batch")
9789 .expect("Skipper must consume all events bytes using writer field types");
9790 assert_eq!(batch.num_rows(), 1);
9791 let ts = batch
9792 .column(0)
9793 .as_any()
9794 .downcast_ref::<StructArray>()
9795 .unwrap();
9796 let seconds = ts
9797 .column_by_name("seconds")
9798 .unwrap()
9799 .as_any()
9800 .downcast_ref::<Int64Array>()
9801 .unwrap();
9802 assert_eq!(seconds.value(0), 100);
9803 }
9804}