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