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