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