1use arrow_array::cast::AsArray;
21use arrow_array::{BooleanArray, RecordBatch, RecordBatchReader};
22use arrow_buffer::{BooleanBuffer, BooleanBufferBuilder};
23use arrow_schema::{ArrowError, DataType as ArrowType, FieldRef, Schema, SchemaRef};
24use arrow_select::filter::filter_record_batch;
25pub use filter::{ArrowPredicate, ArrowPredicateFn, RowFilter};
26use selection::MaskCursor;
27pub use selection::{
28 MaskRunIter, RowSelection, RowSelectionCursor, RowSelectionPolicy, RowSelector,
29};
30use std::fmt::{Debug, Formatter};
31use std::sync::Arc;
32
33pub use crate::arrow::array_reader::RowGroups;
34use crate::arrow::array_reader::{ArrayReader, ArrayReaderBuilder};
35use crate::arrow::schema::{
36 ParquetField, parquet_to_arrow_schema_and_fields, virtual_type::is_virtual_column,
37};
38use crate::arrow::{FieldLevels, ProjectionMask, parquet_to_arrow_field_levels_with_virtual};
39use crate::basic::{BloomFilterAlgorithm, BloomFilterCompression, BloomFilterHash};
40use crate::bloom_filter::{
41 SBBF_HEADER_SIZE_ESTIMATE, Sbbf, chunk_read_bloom_filter_header_and_offset,
42};
43use crate::column::page::{PageIterator, PageReader};
44#[cfg(feature = "encryption")]
45use crate::encryption::decrypt::FileDecryptionProperties;
46use crate::errors::{ParquetError, Result};
47use crate::file::metadata::{
48 PageIndexPolicy, ParquetMetaData, ParquetMetaDataOptions, ParquetMetaDataReader,
49 ParquetStatisticsPolicy, RowGroupMetaData,
50};
51use crate::file::reader::{ChunkReader, SerializedPageReader};
52use crate::schema::types::SchemaDescriptor;
53
54use crate::arrow::arrow_reader::metrics::ArrowReaderMetrics;
55pub use read_plan::{PredicateOptions, ReadPlan, ReadPlanBuilder};
57
58mod filter;
59pub mod metrics;
60mod read_plan;
61pub(crate) mod selection;
62pub mod statistics;
63
64pub const DEFAULT_BATCH_SIZE: usize = 1024;
66
67pub struct ArrowReaderBuilder<T> {
115 pub(crate) input: T,
124
125 pub(crate) metadata: Arc<ParquetMetaData>,
126
127 pub(crate) schema: SchemaRef,
128
129 pub(crate) fields: Option<Arc<ParquetField>>,
130
131 pub(crate) batch_size: usize,
132
133 pub(crate) row_groups: Option<Vec<usize>>,
134
135 pub(crate) projection: ProjectionMask,
136
137 pub(crate) filter: Option<RowFilter>,
138
139 pub(crate) selection: Option<RowSelection>,
140
141 pub(crate) row_selection_policy: RowSelectionPolicy,
142
143 pub(crate) limit: Option<usize>,
144
145 pub(crate) offset: Option<usize>,
146
147 pub(crate) metrics: ArrowReaderMetrics,
148
149 pub(crate) max_predicate_cache_size: usize,
150}
151
152impl<T: Debug> Debug for ArrowReaderBuilder<T> {
153 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
154 f.debug_struct("ArrowReaderBuilder<T>")
155 .field("input", &self.input)
156 .field("metadata", &self.metadata)
157 .field("schema", &self.schema)
158 .field("fields", &self.fields)
159 .field("batch_size", &self.batch_size)
160 .field("row_groups", &self.row_groups)
161 .field("projection", &self.projection)
162 .field("filter", &self.filter)
163 .field("selection", &self.selection)
164 .field("row_selection_policy", &self.row_selection_policy)
165 .field("limit", &self.limit)
166 .field("offset", &self.offset)
167 .field("metrics", &self.metrics)
168 .finish()
169 }
170}
171
172impl<T> ArrowReaderBuilder<T> {
173 pub(crate) fn new_builder(input: T, metadata: ArrowReaderMetadata) -> Self {
174 Self {
175 input,
176 metadata: metadata.metadata,
177 schema: metadata.schema,
178 fields: metadata.fields,
179 batch_size: DEFAULT_BATCH_SIZE,
180 row_groups: None,
181 projection: ProjectionMask::all(),
182 filter: None,
183 selection: None,
184 row_selection_policy: RowSelectionPolicy::default(),
185 limit: None,
186 offset: None,
187 metrics: ArrowReaderMetrics::Disabled,
188 max_predicate_cache_size: 100 * 1024 * 1024, }
190 }
191
192 pub fn metadata(&self) -> &Arc<ParquetMetaData> {
194 &self.metadata
195 }
196
197 pub fn parquet_schema(&self) -> &SchemaDescriptor {
199 self.metadata.file_metadata().schema_descr()
200 }
201
202 pub fn schema(&self) -> &SchemaRef {
204 &self.schema
205 }
206
207 pub fn with_batch_size(self, batch_size: usize) -> Self {
214 let batch_size = batch_size.min(self.metadata.file_metadata().num_rows() as usize);
216 Self { batch_size, ..self }
217 }
218
219 pub fn with_row_groups(self, row_groups: Vec<usize>) -> Self {
223 Self {
224 row_groups: Some(row_groups),
225 ..self
226 }
227 }
228
229 pub fn with_projection(self, mask: ProjectionMask) -> Self {
231 Self {
232 projection: mask,
233 ..self
234 }
235 }
236
237 pub fn with_row_selection_policy(self, policy: RowSelectionPolicy) -> Self {
241 Self {
242 row_selection_policy: policy,
243 ..self
244 }
245 }
246
247 pub fn with_row_selection(self, selection: RowSelection) -> Self {
307 Self {
308 selection: Some(selection),
309 ..self
310 }
311 }
312
313 pub fn with_row_filter(self, filter: RowFilter) -> Self {
353 Self {
354 filter: Some(filter),
355 ..self
356 }
357 }
358
359 pub fn with_limit(self, limit: usize) -> Self {
367 Self {
368 limit: Some(limit),
369 ..self
370 }
371 }
372
373 pub fn with_offset(self, offset: usize) -> Self {
381 Self {
382 offset: Some(offset),
383 ..self
384 }
385 }
386
387 pub fn with_metrics(self, metrics: ArrowReaderMetrics) -> Self {
422 Self { metrics, ..self }
423 }
424
425 pub fn with_max_predicate_cache_size(self, max_predicate_cache_size: usize) -> Self {
440 Self {
441 max_predicate_cache_size,
442 ..self
443 }
444 }
445}
446
447#[derive(Debug, Clone, Default)]
462pub struct ArrowReaderOptions {
463 skip_arrow_metadata: bool,
465 supplied_schema: Option<SchemaRef>,
470
471 pub(crate) column_index: PageIndexPolicy,
472 pub(crate) offset_index: PageIndexPolicy,
473
474 metadata_options: ParquetMetaDataOptions,
476 #[cfg(feature = "encryption")]
478 pub(crate) file_decryption_properties: Option<Arc<FileDecryptionProperties>>,
479
480 virtual_columns: Vec<FieldRef>,
481}
482
483impl ArrowReaderOptions {
484 pub fn new() -> Self {
486 Self::default()
487 }
488
489 pub fn with_skip_arrow_metadata(self, skip_arrow_metadata: bool) -> Self {
496 Self {
497 skip_arrow_metadata,
498 ..self
499 }
500 }
501
502 pub fn with_schema(self, schema: SchemaRef) -> Self {
611 Self {
612 supplied_schema: Some(schema),
613 skip_arrow_metadata: true,
614 ..self
615 }
616 }
617
618 #[deprecated(since = "57.2.0", note = "Use `with_page_index_policy` instead")]
619 pub fn with_page_index(self, page_index: bool) -> Self {
632 self.with_page_index_policy(PageIndexPolicy::from(page_index))
633 }
634
635 pub fn with_page_index_policy(self, policy: PageIndexPolicy) -> Self {
643 self.with_column_index_policy(policy)
644 .with_offset_index_policy(policy)
645 }
646
647 pub fn with_column_index_policy(mut self, policy: PageIndexPolicy) -> Self {
654 self.column_index = policy;
655 self
656 }
657
658 pub fn with_offset_index_policy(mut self, policy: PageIndexPolicy) -> Self {
665 self.offset_index = policy;
666 self
667 }
668
669 pub fn with_parquet_schema(mut self, schema: Arc<SchemaDescriptor>) -> Self {
675 self.metadata_options.set_schema(schema);
676 self
677 }
678
679 pub fn with_encoding_stats_as_mask(mut self, val: bool) -> Self {
690 self.metadata_options.set_encoding_stats_as_mask(val);
691 self
692 }
693
694 pub fn with_encoding_stats_policy(mut self, policy: ParquetStatisticsPolicy) -> Self {
699 self.metadata_options.set_encoding_stats_policy(policy);
700 self
701 }
702
703 pub fn with_column_stats_policy(mut self, policy: ParquetStatisticsPolicy) -> Self {
708 self.metadata_options.set_column_stats_policy(policy);
709 self
710 }
711
712 pub fn with_size_stats_policy(mut self, policy: ParquetStatisticsPolicy) -> Self {
717 self.metadata_options.set_size_stats_policy(policy);
718 self
719 }
720
721 #[cfg(feature = "encryption")]
725 pub fn with_file_decryption_properties(
726 self,
727 file_decryption_properties: Arc<FileDecryptionProperties>,
728 ) -> Self {
729 Self {
730 file_decryption_properties: Some(file_decryption_properties),
731 ..self
732 }
733 }
734
735 pub fn with_virtual_columns(self, virtual_columns: Vec<FieldRef>) -> Result<Self> {
788 for field in &virtual_columns {
790 if !is_virtual_column(field) {
791 return Err(ParquetError::General(format!(
792 "Field '{}' is not a virtual column. Virtual columns must have extension type names starting with 'arrow.virtual.'",
793 field.name()
794 )));
795 }
796 }
797 Ok(Self {
798 virtual_columns,
799 ..self
800 })
801 }
802
803 #[deprecated(
804 since = "57.2.0",
805 note = "Use `column_index_policy` or `offset_index_policy` instead"
806 )]
807 pub fn page_index(&self) -> bool {
814 self.offset_index != PageIndexPolicy::Skip && self.column_index != PageIndexPolicy::Skip
815 }
816
817 pub fn offset_index_policy(&self) -> PageIndexPolicy {
822 self.offset_index
823 }
824
825 pub fn column_index_policy(&self) -> PageIndexPolicy {
830 self.column_index
831 }
832
833 pub fn metadata_options(&self) -> &ParquetMetaDataOptions {
835 &self.metadata_options
836 }
837
838 #[cfg(feature = "encryption")]
843 pub fn file_decryption_properties(&self) -> Option<&Arc<FileDecryptionProperties>> {
844 self.file_decryption_properties.as_ref()
845 }
846}
847
848impl ParquetMetaDataReader {
849 pub fn with_arrow_reader_options(mut self, options: Option<&ArrowReaderOptions>) -> Self {
863 let Some(options) = options else { return self };
864
865 self = self.with_metadata_options(Some(options.metadata_options().clone()));
866
867 #[cfg(feature = "encryption")]
868 {
869 self = self.with_decryption_properties(
870 options.file_decryption_properties.as_ref().map(Arc::clone),
871 );
872 }
873
874 if options.column_index_policy() != PageIndexPolicy::Skip
875 || options.offset_index_policy() != PageIndexPolicy::Skip
876 {
877 self = self
878 .with_column_index_policy(options.column_index_policy())
879 .with_offset_index_policy(options.offset_index_policy());
880 }
881
882 self
883 }
884}
885
886#[derive(Debug, Clone)]
901pub struct ArrowReaderMetadata {
902 pub(crate) metadata: Arc<ParquetMetaData>,
904 pub(crate) schema: SchemaRef,
906 pub(crate) fields: Option<Arc<ParquetField>>,
908}
909
910impl ArrowReaderMetadata {
911 pub fn load<T: ChunkReader>(reader: &T, options: ArrowReaderOptions) -> Result<Self> {
923 let metadata = ParquetMetaDataReader::new()
924 .with_column_index_policy(options.column_index)
925 .with_offset_index_policy(options.offset_index)
926 .with_metadata_options(Some(options.metadata_options.clone()));
927 #[cfg(feature = "encryption")]
928 let metadata = metadata.with_decryption_properties(
929 options.file_decryption_properties.as_ref().map(Arc::clone),
930 );
931 let metadata = metadata.parse_and_finish(reader)?;
932 Self::try_new(Arc::new(metadata), options)
933 }
934
935 pub fn try_new(metadata: Arc<ParquetMetaData>, options: ArrowReaderOptions) -> Result<Self> {
943 match options.supplied_schema {
944 Some(supplied_schema) => Self::with_supplied_schema(
945 metadata,
946 supplied_schema.clone(),
947 &options.virtual_columns,
948 ),
949 None => {
950 let kv_metadata = match options.skip_arrow_metadata {
951 true => None,
952 false => metadata.file_metadata().key_value_metadata(),
953 };
954
955 let (schema, fields) = parquet_to_arrow_schema_and_fields(
956 metadata.file_metadata().schema_descr(),
957 ProjectionMask::all(),
958 kv_metadata,
959 &options.virtual_columns,
960 )?;
961
962 Ok(Self {
963 metadata,
964 schema: Arc::new(schema),
965 fields: fields.map(Arc::new),
966 })
967 }
968 }
969 }
970
971 fn with_supplied_schema(
972 metadata: Arc<ParquetMetaData>,
973 supplied_schema: SchemaRef,
974 virtual_columns: &[FieldRef],
975 ) -> Result<Self> {
976 let parquet_schema = metadata.file_metadata().schema_descr();
977 let field_levels = parquet_to_arrow_field_levels_with_virtual(
978 parquet_schema,
979 ProjectionMask::all(),
980 Some(supplied_schema.fields()),
981 virtual_columns,
982 )?;
983 let fields = field_levels.fields;
984 let inferred_len = fields.len();
985 let supplied_len = supplied_schema.fields().len() + virtual_columns.len();
986 if inferred_len != supplied_len {
990 return Err(arrow_err!(format!(
991 "Incompatible supplied Arrow schema: expected {} columns received {}",
992 inferred_len, supplied_len
993 )));
994 }
995
996 let mut errors = Vec::new();
997
998 let field_iter = supplied_schema.fields().iter().zip(fields.iter());
999
1000 for (field1, field2) in field_iter {
1001 if field1.data_type() != field2.data_type() {
1002 errors.push(format!(
1003 "data type mismatch for field {}: requested {} but found {}",
1004 field1.name(),
1005 field1.data_type(),
1006 field2.data_type()
1007 ));
1008 }
1009 if field1.is_nullable() != field2.is_nullable() {
1010 errors.push(format!(
1011 "nullability mismatch for field {}: expected {:?} but found {:?}",
1012 field1.name(),
1013 field1.is_nullable(),
1014 field2.is_nullable()
1015 ));
1016 }
1017 if field1.metadata() != field2.metadata() {
1018 errors.push(format!(
1019 "metadata mismatch for field {}: expected {:?} but found {:?}",
1020 field1.name(),
1021 field1.metadata(),
1022 field2.metadata()
1023 ));
1024 }
1025 }
1026
1027 if !errors.is_empty() {
1028 let message = errors.join(", ");
1029 return Err(ParquetError::ArrowError(format!(
1030 "Incompatible supplied Arrow schema: {message}",
1031 )));
1032 }
1033
1034 Ok(Self {
1035 metadata,
1036 schema: supplied_schema,
1037 fields: field_levels.levels.map(Arc::new),
1038 })
1039 }
1040
1041 pub fn metadata(&self) -> &Arc<ParquetMetaData> {
1043 &self.metadata
1044 }
1045
1046 pub fn parquet_schema(&self) -> &SchemaDescriptor {
1048 self.metadata.file_metadata().schema_descr()
1049 }
1050
1051 pub fn schema(&self) -> &SchemaRef {
1053 &self.schema
1054 }
1055}
1056
1057#[doc(hidden)]
1058pub struct SyncReader<T: ChunkReader>(T);
1060
1061impl<T: Debug + ChunkReader> Debug for SyncReader<T> {
1062 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
1063 f.debug_tuple("SyncReader").field(&self.0).finish()
1064 }
1065}
1066
1067pub type ParquetRecordBatchReaderBuilder<T> = ArrowReaderBuilder<SyncReader<T>>;
1074
1075impl<T: ChunkReader + 'static> ParquetRecordBatchReaderBuilder<T> {
1076 pub fn try_new(reader: T) -> Result<Self> {
1107 Self::try_new_with_options(reader, Default::default())
1108 }
1109
1110 pub fn try_new_with_options(reader: T, options: ArrowReaderOptions) -> Result<Self> {
1115 let metadata = ArrowReaderMetadata::load(&reader, options)?;
1116 Ok(Self::new_with_metadata(reader, metadata))
1117 }
1118
1119 pub fn new_with_metadata(input: T, metadata: ArrowReaderMetadata) -> Self {
1158 Self::new_builder(SyncReader(input), metadata)
1159 }
1160
1161 pub fn get_row_group_column_bloom_filter(
1167 &self,
1168 row_group_idx: usize,
1169 column_idx: usize,
1170 ) -> Result<Option<Sbbf>> {
1171 let metadata = self.metadata.row_group(row_group_idx);
1172 let column_metadata = metadata.column(column_idx);
1173
1174 let offset: u64 = if let Some(offset) = column_metadata.bloom_filter_offset() {
1175 offset
1176 .try_into()
1177 .map_err(|_| ParquetError::General("Bloom filter offset is invalid".to_string()))?
1178 } else {
1179 return Ok(None);
1180 };
1181
1182 let buffer = match column_metadata.bloom_filter_length() {
1183 Some(length) => self.input.0.get_bytes(offset, length as usize),
1184 None => self.input.0.get_bytes(offset, SBBF_HEADER_SIZE_ESTIMATE),
1185 }?;
1186
1187 let (header, bitset_offset) =
1188 chunk_read_bloom_filter_header_and_offset(offset, buffer.clone())?;
1189
1190 match header.algorithm {
1191 BloomFilterAlgorithm::BLOCK => {
1192 }
1194 }
1195 match header.compression {
1196 BloomFilterCompression::UNCOMPRESSED => {
1197 }
1199 }
1200 match header.hash {
1201 BloomFilterHash::XXHASH => {
1202 }
1204 }
1205
1206 let bitset = match column_metadata.bloom_filter_length() {
1207 Some(_) => buffer.slice(
1208 (TryInto::<usize>::try_into(bitset_offset).unwrap()
1209 - TryInto::<usize>::try_into(offset).unwrap())..,
1210 ),
1211 None => {
1212 let bitset_length: usize = header.num_bytes.try_into().map_err(|_| {
1213 ParquetError::General("Bloom filter length is invalid".to_string())
1214 })?;
1215 self.input.0.get_bytes(bitset_offset, bitset_length)?
1216 }
1217 };
1218 Ok(Some(Sbbf::new(&bitset)))
1219 }
1220
1221 pub fn build(self) -> Result<ParquetRecordBatchReader> {
1225 let Self {
1226 input,
1227 metadata,
1228 schema: _,
1229 fields,
1230 batch_size,
1231 row_groups,
1232 projection,
1233 mut filter,
1234 selection,
1235 row_selection_policy,
1236 limit,
1237 offset,
1238 metrics,
1239 max_predicate_cache_size: _,
1241 } = self;
1242
1243 let batch_size = batch_size.min(metadata.file_metadata().num_rows() as usize);
1245
1246 let row_groups = row_groups.unwrap_or_else(|| (0..metadata.num_row_groups()).collect());
1247
1248 let reader = ReaderRowGroups {
1249 reader: Arc::new(input.0),
1250 metadata,
1251 row_groups,
1252 };
1253
1254 let mut plan_builder = ReadPlanBuilder::new(batch_size)
1255 .with_selection(selection)
1256 .with_row_selection_policy(row_selection_policy);
1257
1258 if let Some(filter) = filter.as_mut() {
1260 for predicate in filter.predicates.iter_mut() {
1261 if !plan_builder.selects_any() {
1263 break;
1264 }
1265
1266 let mut cache_projection = predicate.projection().clone();
1267 cache_projection.intersect(&projection);
1268
1269 let array_reader = ArrayReaderBuilder::new(&reader, &metrics)
1270 .with_batch_size(batch_size)
1271 .with_parquet_metadata(&reader.metadata)
1272 .build_array_reader(fields.as_deref(), predicate.projection())?;
1273
1274 plan_builder = plan_builder.with_predicate(array_reader, predicate.as_mut())?;
1275 }
1276 }
1277
1278 let array_reader = ArrayReaderBuilder::new(&reader, &metrics)
1279 .with_batch_size(batch_size)
1280 .with_parquet_metadata(&reader.metadata)
1281 .build_array_reader(fields.as_deref(), &projection)?;
1282
1283 let read_plan = plan_builder
1284 .limited(reader.num_rows())
1285 .with_offset(offset)
1286 .with_limit(limit)
1287 .build_limited()
1288 .build();
1289
1290 Ok(ParquetRecordBatchReader::new(array_reader, read_plan))
1291 }
1292}
1293
1294struct ReaderRowGroups<T: ChunkReader> {
1295 reader: Arc<T>,
1296
1297 metadata: Arc<ParquetMetaData>,
1298 row_groups: Vec<usize>,
1300}
1301
1302impl<T: ChunkReader + 'static> RowGroups for ReaderRowGroups<T> {
1303 fn num_rows(&self) -> usize {
1304 let meta = self.metadata.row_groups();
1305 self.row_groups
1306 .iter()
1307 .map(|x| meta[*x].num_rows() as usize)
1308 .sum()
1309 }
1310
1311 fn column_chunks(&self, i: usize) -> Result<Box<dyn PageIterator>> {
1312 Ok(Box::new(ReaderPageIterator {
1313 column_idx: i,
1314 reader: self.reader.clone(),
1315 metadata: self.metadata.clone(),
1316 row_groups: self.row_groups.clone().into_iter(),
1317 }))
1318 }
1319
1320 fn row_groups(&self) -> Box<dyn Iterator<Item = &RowGroupMetaData> + '_> {
1321 Box::new(
1322 self.row_groups
1323 .iter()
1324 .map(move |i| self.metadata.row_group(*i)),
1325 )
1326 }
1327
1328 fn metadata(&self) -> &ParquetMetaData {
1329 self.metadata.as_ref()
1330 }
1331}
1332
1333struct ReaderPageIterator<T: ChunkReader> {
1334 reader: Arc<T>,
1335 column_idx: usize,
1336 row_groups: std::vec::IntoIter<usize>,
1337 metadata: Arc<ParquetMetaData>,
1338}
1339
1340impl<T: ChunkReader + 'static> ReaderPageIterator<T> {
1341 fn next_page_reader(&mut self, rg_idx: usize) -> Result<SerializedPageReader<T>> {
1343 let rg = self.metadata.row_group(rg_idx);
1344 let column_chunk_metadata = rg.column(self.column_idx);
1345 let offset_index = self.metadata.offset_index();
1346 let page_locations = offset_index
1349 .filter(|i| !i[rg_idx].is_empty())
1350 .map(|i| i[rg_idx][self.column_idx].page_locations.clone());
1351 let total_rows = rg.num_rows() as usize;
1352 let reader = self.reader.clone();
1353
1354 SerializedPageReader::new(reader, column_chunk_metadata, total_rows, page_locations)?
1355 .add_crypto_context(
1356 rg_idx,
1357 self.column_idx,
1358 self.metadata.as_ref(),
1359 column_chunk_metadata,
1360 )
1361 }
1362}
1363
1364impl<T: ChunkReader + 'static> Iterator for ReaderPageIterator<T> {
1365 type Item = Result<Box<dyn PageReader>>;
1366
1367 fn next(&mut self) -> Option<Self::Item> {
1368 let rg_idx = self.row_groups.next()?;
1369 let page_reader = self
1370 .next_page_reader(rg_idx)
1371 .map(|page_reader| Box::new(page_reader) as _);
1372 Some(page_reader)
1373 }
1374}
1375
1376impl<T: ChunkReader + 'static> PageIterator for ReaderPageIterator<T> {}
1377
1378pub struct ParquetRecordBatchReader {
1389 array_reader: Box<dyn ArrayReader>,
1390 schema: SchemaRef,
1391 read_plan: ReadPlan,
1392}
1393
1394#[derive(Default)]
1417enum FilterMaskAccumulator {
1418 #[default]
1419 Empty,
1420 Single(BooleanBuffer),
1421 Combined(BooleanBufferBuilder),
1422}
1423
1424impl FilterMaskAccumulator {
1425 fn append(&mut self, mask: BooleanBuffer) {
1426 *self = match std::mem::take(self) {
1427 Self::Empty => Self::Single(mask),
1428 Self::Single(first) => {
1429 let mut combined = BooleanBufferBuilder::new(first.len() + mask.len());
1430 combined.append_buffer(&first);
1431 combined.append_buffer(&mask);
1432 Self::Combined(combined)
1433 }
1434 Self::Combined(mut combined) => {
1435 combined.append_buffer(&mask);
1436 Self::Combined(combined)
1437 }
1438 };
1439 }
1440
1441 fn finish(self) -> Option<BooleanBuffer> {
1442 match self {
1443 Self::Empty => None,
1444 Self::Single(mask) => Some(mask),
1445 Self::Combined(combined) => Some(combined.build()),
1446 }
1447 }
1448}
1449
1450fn consume_record_batch(array_reader: &mut dyn ArrayReader) -> Result<RecordBatch> {
1452 let array = array_reader.consume_batch()?;
1453 let struct_array = array.as_struct_opt().ok_or_else(|| {
1454 ArrowError::ParquetError("Struct array reader should return struct array".to_string())
1455 })?;
1456 Ok(RecordBatch::from(struct_array))
1457}
1458
1459fn read_mask_batch(
1467 array_reader: &mut dyn ArrayReader,
1468 mask_cursor: &mut MaskCursor,
1469 batch_size: usize,
1470) -> Result<Option<RecordBatch>> {
1471 let mut selected_rows = 0;
1472 let mut filter_mask = FilterMaskAccumulator::default();
1473
1474 while selected_rows < batch_size && !mask_cursor.is_empty() {
1475 let mask_chunk = mask_cursor.next_chunk(batch_size - selected_rows)?;
1476
1477 if mask_chunk.initial_skip > 0 {
1478 let skipped = array_reader.skip_records(mask_chunk.initial_skip)?;
1479 if skipped != mask_chunk.initial_skip {
1480 return Err(general_err!(
1481 "failed to skip rows, expected {}, got {}",
1482 mask_chunk.initial_skip,
1483 skipped
1484 ));
1485 }
1486 }
1487
1488 let mask = mask_cursor.mask_values_for(&mask_chunk)?;
1489 let read = array_reader.read_records(mask_chunk.chunk_rows)?;
1490 if read == 0 {
1491 return Err(general_err!(
1492 "reached end of column while expecting {} rows",
1493 mask_chunk.chunk_rows
1494 ));
1495 }
1496 if read != mask_chunk.chunk_rows {
1497 return Err(general_err!(
1498 "insufficient rows read from array reader - expected {}, got {}",
1499 mask_chunk.chunk_rows,
1500 read
1501 ));
1502 }
1503
1504 filter_mask.append(mask.values().clone());
1505 selected_rows += mask_chunk.selected_rows;
1506 }
1507
1508 if selected_rows == 0 {
1509 return Ok(None);
1510 }
1511
1512 let filter_mask = filter_mask
1513 .finish()
1514 .ok_or_else(|| general_err!("Internal Error: decoded Mask batch has no filter values"))?;
1515 let batch = consume_record_batch(array_reader)?;
1516 let filtered_batch = filter_record_batch(&batch, &BooleanArray::from(filter_mask))?;
1517 if filtered_batch.num_rows() != selected_rows {
1518 return Err(general_err!(
1519 "filtered rows mismatch selection - expected {}, got {}",
1520 selected_rows,
1521 filtered_batch.num_rows()
1522 ));
1523 }
1524
1525 Ok(Some(filtered_batch))
1526}
1527
1528impl Debug for ParquetRecordBatchReader {
1529 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
1530 f.debug_struct("ParquetRecordBatchReader")
1531 .field("array_reader", &"...")
1532 .field("schema", &self.schema)
1533 .field("read_plan", &self.read_plan)
1534 .finish()
1535 }
1536}
1537
1538impl Iterator for ParquetRecordBatchReader {
1539 type Item = Result<RecordBatch, ArrowError>;
1540
1541 fn next(&mut self) -> Option<Self::Item> {
1542 self.next_inner()
1543 .map_err(|arrow_err| arrow_err.into())
1544 .transpose()
1545 }
1546}
1547
1548impl ParquetRecordBatchReader {
1549 fn next_inner(&mut self) -> Result<Option<RecordBatch>> {
1555 let mut read_records = 0;
1556 let batch_size = self.batch_size();
1557 if batch_size == 0 {
1558 return Ok(None);
1559 }
1560 match self.read_plan.row_selection_cursor_mut() {
1561 RowSelectionCursor::Mask(mask_cursor) => {
1562 return read_mask_batch(self.array_reader.as_mut(), mask_cursor, batch_size);
1563 }
1564 RowSelectionCursor::Selectors(selectors_cursor) => {
1565 while read_records < batch_size && !selectors_cursor.is_empty() {
1566 let front = selectors_cursor.next_selector();
1567 if front.skip {
1568 let skipped = self.array_reader.skip_records(front.row_count)?;
1569
1570 if skipped != front.row_count {
1571 return Err(general_err!(
1572 "failed to skip rows, expected {}, got {}",
1573 front.row_count,
1574 skipped
1575 ));
1576 }
1577 continue;
1578 }
1579
1580 if front.row_count == 0 {
1583 continue;
1584 }
1585
1586 let need_read = batch_size - read_records;
1588 let to_read = match front.row_count.checked_sub(need_read) {
1589 Some(remaining) if remaining != 0 => {
1590 selectors_cursor.return_selector(RowSelector::select(remaining));
1593 need_read
1594 }
1595 _ => front.row_count,
1596 };
1597 match self.array_reader.read_records(to_read)? {
1598 0 => break,
1599 rec => read_records += rec,
1600 };
1601 }
1602 }
1603 RowSelectionCursor::All => {
1604 self.array_reader.read_records(batch_size)?;
1605 }
1606 }
1607
1608 let batch = consume_record_batch(self.array_reader.as_mut())?;
1609 Ok(if batch.num_rows() > 0 {
1610 Some(batch)
1611 } else {
1612 None
1613 })
1614 }
1615}
1616
1617impl RecordBatchReader for ParquetRecordBatchReader {
1618 fn schema(&self) -> SchemaRef {
1623 self.schema.clone()
1624 }
1625}
1626
1627impl ParquetRecordBatchReader {
1628 pub fn try_new<T: ChunkReader + 'static>(reader: T, batch_size: usize) -> Result<Self> {
1632 ParquetRecordBatchReaderBuilder::try_new(reader)?
1633 .with_batch_size(batch_size)
1634 .build()
1635 }
1636
1637 pub fn try_new_with_row_groups(
1642 levels: &FieldLevels,
1643 row_groups: &dyn RowGroups,
1644 batch_size: usize,
1645 selection: Option<RowSelection>,
1646 ) -> Result<Self> {
1647 let metrics = ArrowReaderMetrics::disabled();
1649 let array_reader = ArrayReaderBuilder::new(row_groups, &metrics)
1650 .with_batch_size(batch_size)
1651 .with_parquet_metadata(row_groups.metadata())
1652 .build_array_reader(levels.levels.as_ref(), &ProjectionMask::all())?;
1653
1654 let read_plan = ReadPlanBuilder::new(batch_size)
1655 .with_selection(selection)
1656 .build();
1657
1658 Ok(Self {
1659 array_reader,
1660 schema: Arc::new(Schema::new(levels.fields.clone())),
1661 read_plan,
1662 })
1663 }
1664
1665 pub(crate) fn new(array_reader: Box<dyn ArrayReader>, read_plan: ReadPlan) -> Self {
1669 let schema = match array_reader.get_data_type() {
1670 ArrowType::Struct(fields) => Schema::new(fields.clone()),
1671 _ => unreachable!("Struct array reader's data type is not struct!"),
1672 };
1673
1674 Self {
1675 array_reader,
1676 schema: Arc::new(schema),
1677 read_plan,
1678 }
1679 }
1680
1681 #[inline(always)]
1682 pub(crate) fn batch_size(&self) -> usize {
1683 self.read_plan.batch_size()
1684 }
1685}
1686
1687#[cfg(test)]
1688pub(crate) mod tests {
1689 use std::cmp::min;
1690 use std::collections::{HashMap, VecDeque};
1691 use std::fmt::Formatter;
1692 use std::fs::File;
1693 use std::io::Seek;
1694 use std::path::PathBuf;
1695 use std::sync::Arc;
1696
1697 use rand::rngs::StdRng;
1698 use rand::{Rng, RngCore, SeedableRng, random, rng};
1699 use tempfile::tempfile;
1700
1701 use crate::arrow::arrow_reader::{
1702 ArrowPredicateFn, ArrowReaderMetadata, ArrowReaderOptions, ParquetRecordBatchReader,
1703 ParquetRecordBatchReaderBuilder, RowFilter, RowSelection, RowSelector,
1704 };
1705 use crate::arrow::schema::{
1706 add_encoded_arrow_schema_to_metadata,
1707 virtual_type::{RowGroupIndex, RowNumber},
1708 };
1709 use crate::arrow::{ArrowWriter, ProjectionMask};
1710 use crate::basic::{ConvertedType, Encoding, LogicalType, Repetition, Type as PhysicalType};
1711 use crate::column::reader::decoder::REPETITION_LEVELS_BATCH_SIZE;
1712 use crate::data_type::{
1713 BoolType, ByteArray, ByteArrayType, DataType, DoubleType, FixedLenByteArray,
1714 FixedLenByteArrayType, FloatType, Int32Type, Int64Type, Int96, Int96Type,
1715 };
1716 use crate::errors::Result;
1717 use crate::file::metadata::{PageIndexPolicy, ParquetMetaData, ParquetStatisticsPolicy};
1718 use crate::file::properties::{EnabledStatistics, WriterProperties, WriterVersion};
1719 use crate::file::writer::{SerializedFileWriter, SerializedRowGroupWriter};
1720 use crate::schema::parser::parse_message_type;
1721 use crate::schema::types::{Type, TypePtr};
1722 use crate::util::test_common::rand_gen::RandGen;
1723 use arrow_array::builder::*;
1724 use arrow_array::cast::AsArray;
1725 use arrow_array::types::{
1726 Date32Type, Date64Type, Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type,
1727 DecimalType, Float16Type, Float32Type, Float64Type, Time32MillisecondType,
1728 Time64MicrosecondType,
1729 };
1730 use arrow_array::*;
1731 use arrow_buffer::{ArrowNativeType, BooleanBuffer, Buffer, IntervalDayTime, NullBuffer, i256};
1732 use arrow_data::{ArrayData, ArrayDataBuilder};
1733 use arrow_schema::{DataType as ArrowDataType, Field, Fields, Schema, SchemaRef, TimeUnit};
1734 use arrow_select::concat::concat_batches;
1735 use bytes::Bytes;
1736 use half::f16;
1737 use num_traits::PrimInt;
1738
1739 #[test]
1740 fn filter_mask_accumulator_handles_empty_single_and_multiple_chunks() {
1741 let first = BooleanBuffer::from(vec![true, false, false, false]);
1742 let second = BooleanBuffer::from(vec![true]);
1743 let third = BooleanBuffer::from(vec![false, true]);
1744
1745 assert!(super::FilterMaskAccumulator::default().finish().is_none());
1746
1747 let mut single = super::FilterMaskAccumulator::default();
1748 single.append(first.clone());
1749 assert_eq!(single.finish().unwrap(), first);
1750
1751 let mut combined = super::FilterMaskAccumulator::default();
1752 combined.append(BooleanBuffer::from(vec![true, false, false, false]));
1753 combined.append(second);
1754 combined.append(third);
1755 assert_eq!(
1756 combined.finish().unwrap(),
1757 BooleanBuffer::from(vec![true, false, false, false, true, false, true])
1758 );
1759 }
1760
1761 #[test]
1762 fn test_arrow_reader_all_columns() {
1763 let file = get_test_file("parquet/generated_simple_numerics/blogs.parquet");
1764
1765 let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
1766 let original_schema = Arc::clone(builder.schema());
1767 let reader = builder.build().unwrap();
1768
1769 assert_eq!(original_schema.fields(), reader.schema().fields());
1771 }
1772
1773 #[test]
1774 fn test_reuse_schema() {
1775 let file = get_test_file("parquet/alltypes-java.parquet");
1776
1777 let builder = ParquetRecordBatchReaderBuilder::try_new(file.try_clone().unwrap()).unwrap();
1778 let expected = builder.metadata;
1779 let schema = expected.file_metadata().schema_descr_ptr();
1780
1781 let arrow_options = ArrowReaderOptions::new().with_parquet_schema(schema.clone());
1782 let builder =
1783 ParquetRecordBatchReaderBuilder::try_new_with_options(file, arrow_options).unwrap();
1784
1785 assert_eq!(expected.as_ref(), builder.metadata.as_ref());
1787 }
1788
1789 #[test]
1790 fn test_page_encoding_stats_mask() {
1791 let testdata = arrow::util::test_util::parquet_test_data();
1792 let path = format!("{testdata}/alltypes_tiny_pages.parquet");
1793 let file = File::open(path).unwrap();
1794
1795 let arrow_options = ArrowReaderOptions::new().with_encoding_stats_as_mask(true);
1796 let builder =
1797 ParquetRecordBatchReaderBuilder::try_new_with_options(file, arrow_options).unwrap();
1798
1799 let row_group_metadata = builder.metadata.row_group(0);
1800
1801 let page_encoding_stats = row_group_metadata
1803 .column(0)
1804 .page_encoding_stats_mask()
1805 .unwrap();
1806 assert!(page_encoding_stats.is_only(Encoding::PLAIN));
1807 let page_encoding_stats = row_group_metadata
1808 .column(2)
1809 .page_encoding_stats_mask()
1810 .unwrap();
1811 assert!(page_encoding_stats.is_only(Encoding::PLAIN_DICTIONARY));
1812 }
1813
1814 #[test]
1815 fn test_stats_stats_skipped() {
1816 let testdata = arrow::util::test_util::parquet_test_data();
1817 let path = format!("{testdata}/alltypes_tiny_pages.parquet");
1818 let file = File::open(path).unwrap();
1819
1820 let arrow_options = ArrowReaderOptions::new()
1822 .with_encoding_stats_policy(ParquetStatisticsPolicy::SkipAll)
1823 .with_column_stats_policy(ParquetStatisticsPolicy::SkipAll);
1824 let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(
1825 file.try_clone().unwrap(),
1826 arrow_options,
1827 )
1828 .unwrap();
1829
1830 let row_group_metadata = builder.metadata.row_group(0);
1831 for column in row_group_metadata.columns() {
1832 assert!(column.page_encoding_stats().is_none());
1833 assert!(column.page_encoding_stats_mask().is_none());
1834 assert!(column.statistics().is_none());
1835 }
1836
1837 let arrow_options = ArrowReaderOptions::new()
1839 .with_encoding_stats_as_mask(true)
1840 .with_encoding_stats_policy(ParquetStatisticsPolicy::skip_except(&[0]))
1841 .with_column_stats_policy(ParquetStatisticsPolicy::skip_except(&[0]));
1842 let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(
1843 file.try_clone().unwrap(),
1844 arrow_options,
1845 )
1846 .unwrap();
1847
1848 let row_group_metadata = builder.metadata.row_group(0);
1849 for (idx, column) in row_group_metadata.columns().iter().enumerate() {
1850 assert!(column.page_encoding_stats().is_none());
1851 assert_eq!(column.page_encoding_stats_mask().is_some(), idx == 0);
1852 assert_eq!(column.statistics().is_some(), idx == 0);
1853 }
1854 }
1855
1856 #[test]
1857 fn test_size_stats_stats_skipped() {
1858 let testdata = arrow::util::test_util::parquet_test_data();
1859 let path = format!("{testdata}/repeated_primitive_no_list.parquet");
1860 let file = File::open(path).unwrap();
1861
1862 let arrow_options =
1864 ArrowReaderOptions::new().with_size_stats_policy(ParquetStatisticsPolicy::SkipAll);
1865 let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(
1866 file.try_clone().unwrap(),
1867 arrow_options,
1868 )
1869 .unwrap();
1870
1871 let row_group_metadata = builder.metadata.row_group(0);
1872 for column in row_group_metadata.columns() {
1873 assert!(column.repetition_level_histogram().is_none());
1874 assert!(column.definition_level_histogram().is_none());
1875 assert!(column.unencoded_byte_array_data_bytes().is_none());
1876 }
1877
1878 let arrow_options = ArrowReaderOptions::new()
1880 .with_encoding_stats_as_mask(true)
1881 .with_size_stats_policy(ParquetStatisticsPolicy::skip_except(&[1]));
1882 let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(
1883 file.try_clone().unwrap(),
1884 arrow_options,
1885 )
1886 .unwrap();
1887
1888 let row_group_metadata = builder.metadata.row_group(0);
1889 for (idx, column) in row_group_metadata.columns().iter().enumerate() {
1890 assert_eq!(column.repetition_level_histogram().is_some(), idx == 1);
1891 assert_eq!(column.definition_level_histogram().is_some(), idx == 1);
1892 assert_eq!(column.unencoded_byte_array_data_bytes().is_some(), idx == 1);
1893 }
1894 }
1895
1896 #[test]
1897 fn test_arrow_reader_single_column() {
1898 let file = get_test_file("parquet/generated_simple_numerics/blogs.parquet");
1899
1900 let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
1901 let original_schema = Arc::clone(builder.schema());
1902
1903 let mask = ProjectionMask::leaves(builder.parquet_schema(), [2]);
1904 let reader = builder.with_projection(mask).build().unwrap();
1905
1906 assert_eq!(1, reader.schema().fields().len());
1908 assert_eq!(original_schema.fields()[1], reader.schema().fields()[0]);
1909 }
1910
1911 #[test]
1912 fn test_arrow_reader_single_column_by_name() {
1913 let file = get_test_file("parquet/generated_simple_numerics/blogs.parquet");
1914
1915 let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
1916 let original_schema = Arc::clone(builder.schema());
1917
1918 let mask = ProjectionMask::columns(builder.parquet_schema(), ["blog_id"]);
1919 let reader = builder.with_projection(mask).build().unwrap();
1920
1921 assert_eq!(1, reader.schema().fields().len());
1923 assert_eq!(original_schema.fields()[1], reader.schema().fields()[0]);
1924 }
1925
1926 #[test]
1927 fn test_null_column_reader_test() {
1928 let mut file = tempfile::tempfile().unwrap();
1929
1930 let schema = "
1931 message message {
1932 OPTIONAL INT32 int32;
1933 }
1934 ";
1935 let schema = Arc::new(parse_message_type(schema).unwrap());
1936
1937 let def_levels = vec![vec![0, 0, 0], vec![0, 0, 0, 0]];
1938 generate_single_column_file_with_data::<Int32Type>(
1939 &[vec![], vec![]],
1940 Some(&def_levels),
1941 file.try_clone().unwrap(), schema,
1943 Some(Field::new("int32", ArrowDataType::Null, true)),
1944 &Default::default(),
1945 )
1946 .unwrap();
1947
1948 file.rewind().unwrap();
1949
1950 let record_reader = ParquetRecordBatchReader::try_new(file, 2).unwrap();
1951 let batches = record_reader.collect::<Result<Vec<_>, _>>().unwrap();
1952
1953 assert_eq!(batches.len(), 4);
1954 for batch in &batches[0..3] {
1955 assert_eq!(batch.num_rows(), 2);
1956 assert_eq!(batch.num_columns(), 1);
1957 assert_eq!(batch.column(0).null_count(), 2);
1958 }
1959
1960 assert_eq!(batches[3].num_rows(), 1);
1961 assert_eq!(batches[3].num_columns(), 1);
1962 assert_eq!(batches[3].column(0).null_count(), 1);
1963 }
1964
1965 #[test]
1966 fn test_primitive_single_column_reader_test() {
1967 run_single_column_reader_tests::<BoolType, _, BoolType>(
1968 2,
1969 ConvertedType::NONE,
1970 None,
1971 |vals| Arc::new(BooleanArray::from_iter(vals.iter().cloned())),
1972 &[Encoding::PLAIN, Encoding::RLE, Encoding::RLE_DICTIONARY],
1973 );
1974 run_single_column_reader_tests::<Int32Type, _, Int32Type>(
1975 2,
1976 ConvertedType::NONE,
1977 None,
1978 |vals| Arc::new(Int32Array::from_iter(vals.iter().cloned())),
1979 &[
1980 Encoding::PLAIN,
1981 Encoding::RLE_DICTIONARY,
1982 Encoding::DELTA_BINARY_PACKED,
1983 Encoding::BYTE_STREAM_SPLIT,
1984 ],
1985 );
1986 run_single_column_reader_tests::<Int64Type, _, Int64Type>(
1987 2,
1988 ConvertedType::NONE,
1989 None,
1990 |vals| Arc::new(Int64Array::from_iter(vals.iter().cloned())),
1991 &[
1992 Encoding::PLAIN,
1993 Encoding::RLE_DICTIONARY,
1994 Encoding::DELTA_BINARY_PACKED,
1995 Encoding::BYTE_STREAM_SPLIT,
1996 ],
1997 );
1998 run_single_column_reader_tests::<FloatType, _, FloatType>(
1999 2,
2000 ConvertedType::NONE,
2001 None,
2002 |vals| Arc::new(Float32Array::from_iter(vals.iter().cloned())),
2003 &[Encoding::PLAIN, Encoding::BYTE_STREAM_SPLIT],
2004 );
2005 }
2006
2007 #[test]
2008 fn test_unsigned_primitive_single_column_reader_test() {
2009 run_single_column_reader_tests::<Int32Type, _, Int32Type>(
2010 2,
2011 ConvertedType::UINT_32,
2012 Some(ArrowDataType::UInt32),
2013 |vals| {
2014 Arc::new(UInt32Array::from_iter(
2015 vals.iter().map(|x| x.map(|x| x as u32)),
2016 ))
2017 },
2018 &[
2019 Encoding::PLAIN,
2020 Encoding::RLE_DICTIONARY,
2021 Encoding::DELTA_BINARY_PACKED,
2022 ],
2023 );
2024 run_single_column_reader_tests::<Int64Type, _, Int64Type>(
2025 2,
2026 ConvertedType::UINT_64,
2027 Some(ArrowDataType::UInt64),
2028 |vals| {
2029 Arc::new(UInt64Array::from_iter(
2030 vals.iter().map(|x| x.map(|x| x as u64)),
2031 ))
2032 },
2033 &[
2034 Encoding::PLAIN,
2035 Encoding::RLE_DICTIONARY,
2036 Encoding::DELTA_BINARY_PACKED,
2037 ],
2038 );
2039 }
2040
2041 #[test]
2042 fn test_unsigned_roundtrip() {
2043 let schema = Arc::new(Schema::new(vec![
2044 Field::new("uint32", ArrowDataType::UInt32, true),
2045 Field::new("uint64", ArrowDataType::UInt64, true),
2046 ]));
2047
2048 let mut buf = Vec::with_capacity(1024);
2049 let mut writer = ArrowWriter::try_new(&mut buf, schema.clone(), None).unwrap();
2050
2051 let original = RecordBatch::try_new(
2052 schema,
2053 vec![
2054 Arc::new(UInt32Array::from_iter_values([
2055 0,
2056 i32::MAX as u32,
2057 u32::MAX,
2058 ])),
2059 Arc::new(UInt64Array::from_iter_values([
2060 0,
2061 i64::MAX as u64,
2062 u64::MAX,
2063 ])),
2064 ],
2065 )
2066 .unwrap();
2067
2068 writer.write(&original).unwrap();
2069 writer.close().unwrap();
2070
2071 let mut reader = ParquetRecordBatchReader::try_new(Bytes::from(buf), 1024).unwrap();
2072 let ret = reader.next().unwrap().unwrap();
2073 assert_eq!(ret, original);
2074
2075 ret.column(0)
2077 .as_any()
2078 .downcast_ref::<UInt32Array>()
2079 .unwrap();
2080
2081 ret.column(1)
2082 .as_any()
2083 .downcast_ref::<UInt64Array>()
2084 .unwrap();
2085 }
2086
2087 #[test]
2088 fn test_float16_roundtrip() -> Result<()> {
2089 let schema = Arc::new(Schema::new(vec![
2090 Field::new("float16", ArrowDataType::Float16, false),
2091 Field::new("float16-nullable", ArrowDataType::Float16, true),
2092 ]));
2093
2094 let mut buf = Vec::with_capacity(1024);
2095 let mut writer = ArrowWriter::try_new(&mut buf, schema.clone(), None)?;
2096
2097 let original = RecordBatch::try_new(
2098 schema,
2099 vec![
2100 Arc::new(Float16Array::from_iter_values([
2101 f16::EPSILON,
2102 f16::MIN,
2103 f16::MAX,
2104 f16::NAN,
2105 f16::INFINITY,
2106 f16::NEG_INFINITY,
2107 f16::ONE,
2108 f16::NEG_ONE,
2109 f16::ZERO,
2110 f16::NEG_ZERO,
2111 f16::E,
2112 f16::PI,
2113 f16::FRAC_1_PI,
2114 ])),
2115 Arc::new(Float16Array::from(vec![
2116 None,
2117 None,
2118 None,
2119 Some(f16::NAN),
2120 Some(f16::INFINITY),
2121 Some(f16::NEG_INFINITY),
2122 None,
2123 None,
2124 None,
2125 None,
2126 None,
2127 None,
2128 Some(f16::FRAC_1_PI),
2129 ])),
2130 ],
2131 )?;
2132
2133 writer.write(&original)?;
2134 writer.close()?;
2135
2136 let mut reader = ParquetRecordBatchReader::try_new(Bytes::from(buf), 1024)?;
2137 let ret = reader.next().unwrap()?;
2138 assert_eq!(ret, original);
2139
2140 ret.column(0).as_primitive::<Float16Type>();
2142 ret.column(1).as_primitive::<Float16Type>();
2143
2144 Ok(())
2145 }
2146
2147 #[test]
2148 fn test_time_utc_roundtrip() -> Result<()> {
2149 let schema = Arc::new(Schema::new(vec![
2150 Field::new(
2151 "time_millis",
2152 ArrowDataType::Time32(TimeUnit::Millisecond),
2153 true,
2154 )
2155 .with_metadata(HashMap::from_iter(vec![(
2156 "adjusted_to_utc".to_string(),
2157 "".to_string(),
2158 )])),
2159 Field::new(
2160 "time_micros",
2161 ArrowDataType::Time64(TimeUnit::Microsecond),
2162 true,
2163 )
2164 .with_metadata(HashMap::from_iter(vec![(
2165 "adjusted_to_utc".to_string(),
2166 "".to_string(),
2167 )])),
2168 ]));
2169
2170 let mut buf = Vec::with_capacity(1024);
2171 let mut writer = ArrowWriter::try_new(&mut buf, schema.clone(), None)?;
2172
2173 let original = RecordBatch::try_new(
2174 schema,
2175 vec![
2176 Arc::new(Time32MillisecondArray::from(vec![
2177 Some(-1),
2178 Some(0),
2179 Some(86_399_000),
2180 Some(86_400_000),
2181 Some(86_401_000),
2182 None,
2183 ])),
2184 Arc::new(Time64MicrosecondArray::from(vec![
2185 Some(-1),
2186 Some(0),
2187 Some(86_399 * 1_000_000),
2188 Some(86_400 * 1_000_000),
2189 Some(86_401 * 1_000_000),
2190 None,
2191 ])),
2192 ],
2193 )?;
2194
2195 writer.write(&original)?;
2196 writer.close()?;
2197
2198 let mut reader = ParquetRecordBatchReader::try_new(Bytes::from(buf), 1024)?;
2199 let ret = reader.next().unwrap()?;
2200 assert_eq!(ret, original);
2201
2202 ret.column(0).as_primitive::<Time32MillisecondType>();
2204 ret.column(1).as_primitive::<Time64MicrosecondType>();
2205
2206 Ok(())
2207 }
2208
2209 #[test]
2210 fn test_date32_roundtrip() -> Result<()> {
2211 use arrow_array::Date32Array;
2212
2213 let schema = Arc::new(Schema::new(vec![Field::new(
2214 "date32",
2215 ArrowDataType::Date32,
2216 false,
2217 )]));
2218
2219 let mut buf = Vec::with_capacity(1024);
2220
2221 let mut writer = ArrowWriter::try_new(&mut buf, schema.clone(), None)?;
2222
2223 let original = RecordBatch::try_new(
2224 schema,
2225 vec![Arc::new(Date32Array::from(vec![
2226 -1_000_000, -100_000, -10_000, -1_000, 0, 1_000, 10_000, 100_000, 1_000_000,
2227 ]))],
2228 )?;
2229
2230 writer.write(&original)?;
2231 writer.close()?;
2232
2233 let mut reader = ParquetRecordBatchReader::try_new(Bytes::from(buf), 1024)?;
2234 let ret = reader.next().unwrap()?;
2235 assert_eq!(ret, original);
2236
2237 ret.column(0).as_primitive::<Date32Type>();
2239
2240 Ok(())
2241 }
2242
2243 #[test]
2244 fn test_date64_roundtrip() -> Result<()> {
2245 use arrow_array::Date64Array;
2246
2247 let schema = Arc::new(Schema::new(vec![
2248 Field::new("small-date64", ArrowDataType::Date64, false),
2249 Field::new("big-date64", ArrowDataType::Date64, false),
2250 Field::new("invalid-date64", ArrowDataType::Date64, false),
2251 ]));
2252
2253 let mut default_buf = Vec::with_capacity(1024);
2254 let mut coerce_buf = Vec::with_capacity(1024);
2255
2256 let coerce_props = WriterProperties::builder().set_coerce_types(true).build();
2257
2258 let mut default_writer = ArrowWriter::try_new(&mut default_buf, schema.clone(), None)?;
2259 let mut coerce_writer =
2260 ArrowWriter::try_new(&mut coerce_buf, schema.clone(), Some(coerce_props))?;
2261
2262 static NUM_MILLISECONDS_IN_DAY: i64 = 1000 * 60 * 60 * 24;
2263
2264 let original = RecordBatch::try_new(
2265 schema,
2266 vec![
2267 Arc::new(Date64Array::from(vec![
2269 -1_000_000 * NUM_MILLISECONDS_IN_DAY,
2270 -1_000 * NUM_MILLISECONDS_IN_DAY,
2271 0,
2272 1_000 * NUM_MILLISECONDS_IN_DAY,
2273 1_000_000 * NUM_MILLISECONDS_IN_DAY,
2274 ])),
2275 Arc::new(Date64Array::from(vec![
2277 -10_000_000_000 * NUM_MILLISECONDS_IN_DAY,
2278 -1_000_000_000 * NUM_MILLISECONDS_IN_DAY,
2279 0,
2280 1_000_000_000 * NUM_MILLISECONDS_IN_DAY,
2281 10_000_000_000 * NUM_MILLISECONDS_IN_DAY,
2282 ])),
2283 Arc::new(Date64Array::from(vec![
2285 -1_000_000 * NUM_MILLISECONDS_IN_DAY + 1,
2286 -1_000 * NUM_MILLISECONDS_IN_DAY + 1,
2287 1,
2288 1_000 * NUM_MILLISECONDS_IN_DAY + 1,
2289 1_000_000 * NUM_MILLISECONDS_IN_DAY + 1,
2290 ])),
2291 ],
2292 )?;
2293
2294 default_writer.write(&original)?;
2295 coerce_writer.write(&original)?;
2296
2297 default_writer.close()?;
2298 coerce_writer.close()?;
2299
2300 let mut default_reader = ParquetRecordBatchReader::try_new(Bytes::from(default_buf), 1024)?;
2301 let mut coerce_reader = ParquetRecordBatchReader::try_new(Bytes::from(coerce_buf), 1024)?;
2302
2303 let default_ret = default_reader.next().unwrap()?;
2304 let coerce_ret = coerce_reader.next().unwrap()?;
2305
2306 assert_eq!(default_ret, original);
2308
2309 assert_eq!(coerce_ret.column(0), original.column(0));
2311 assert_ne!(coerce_ret.column(1), original.column(1));
2312 assert_ne!(coerce_ret.column(2), original.column(2));
2313
2314 default_ret.column(0).as_primitive::<Date64Type>();
2316 coerce_ret.column(0).as_primitive::<Date64Type>();
2317
2318 Ok(())
2319 }
2320 struct RandFixedLenGen {}
2321
2322 impl RandGen<FixedLenByteArrayType> for RandFixedLenGen {
2323 fn r#gen(len: i32) -> FixedLenByteArray {
2324 let mut v = vec![0u8; len as usize];
2325 rng().fill_bytes(&mut v);
2326 ByteArray::from(v).into()
2327 }
2328 }
2329
2330 #[test]
2331 fn test_fixed_length_binary_column_reader() {
2332 run_single_column_reader_tests::<FixedLenByteArrayType, _, RandFixedLenGen>(
2333 20,
2334 ConvertedType::NONE,
2335 None,
2336 |vals| {
2337 let mut builder = FixedSizeBinaryBuilder::with_capacity(vals.len(), 20);
2338 for val in vals {
2339 match val {
2340 Some(b) => builder.append_value(b).unwrap(),
2341 None => builder.append_null(),
2342 }
2343 }
2344 Arc::new(builder.finish())
2345 },
2346 &[Encoding::PLAIN, Encoding::RLE_DICTIONARY],
2347 );
2348 }
2349
2350 #[test]
2351 fn test_interval_day_time_column_reader() {
2352 run_single_column_reader_tests::<FixedLenByteArrayType, _, RandFixedLenGen>(
2353 12,
2354 ConvertedType::INTERVAL,
2355 None,
2356 |vals| {
2357 Arc::new(
2358 vals.iter()
2359 .map(|x| {
2360 x.as_ref().map(|b| IntervalDayTime {
2361 days: i32::from_le_bytes(b.as_ref()[4..8].try_into().unwrap()),
2362 milliseconds: i32::from_le_bytes(
2363 b.as_ref()[8..12].try_into().unwrap(),
2364 ),
2365 })
2366 })
2367 .collect::<IntervalDayTimeArray>(),
2368 )
2369 },
2370 &[Encoding::PLAIN, Encoding::RLE_DICTIONARY],
2371 );
2372 }
2373
2374 #[test]
2375 fn test_int96_single_column_reader_test() {
2376 let encodings = &[Encoding::PLAIN, Encoding::RLE_DICTIONARY];
2377
2378 type TypeHintAndConversionFunction =
2379 (Option<ArrowDataType>, fn(&[Option<Int96>]) -> ArrayRef);
2380
2381 let resolutions: Vec<TypeHintAndConversionFunction> = vec![
2382 (None, |vals: &[Option<Int96>]| {
2384 Arc::new(TimestampNanosecondArray::from_iter(
2385 vals.iter().map(|x| x.map(|x| x.to_nanos())),
2386 )) as ArrayRef
2387 }),
2388 (
2390 Some(ArrowDataType::Timestamp(TimeUnit::Second, None)),
2391 |vals: &[Option<Int96>]| {
2392 Arc::new(TimestampSecondArray::from_iter(
2393 vals.iter().map(|x| x.map(|x| x.to_seconds())),
2394 )) as ArrayRef
2395 },
2396 ),
2397 (
2398 Some(ArrowDataType::Timestamp(TimeUnit::Millisecond, None)),
2399 |vals: &[Option<Int96>]| {
2400 Arc::new(TimestampMillisecondArray::from_iter(
2401 vals.iter().map(|x| x.map(|x| x.to_millis())),
2402 )) as ArrayRef
2403 },
2404 ),
2405 (
2406 Some(ArrowDataType::Timestamp(TimeUnit::Microsecond, None)),
2407 |vals: &[Option<Int96>]| {
2408 Arc::new(TimestampMicrosecondArray::from_iter(
2409 vals.iter().map(|x| x.map(|x| x.to_micros())),
2410 )) as ArrayRef
2411 },
2412 ),
2413 (
2414 Some(ArrowDataType::Timestamp(TimeUnit::Nanosecond, None)),
2415 |vals: &[Option<Int96>]| {
2416 Arc::new(TimestampNanosecondArray::from_iter(
2417 vals.iter().map(|x| x.map(|x| x.to_nanos())),
2418 )) as ArrayRef
2419 },
2420 ),
2421 (
2423 Some(ArrowDataType::Timestamp(
2424 TimeUnit::Second,
2425 Some(Arc::from("-05:00")),
2426 )),
2427 |vals: &[Option<Int96>]| {
2428 Arc::new(
2429 TimestampSecondArray::from_iter(
2430 vals.iter().map(|x| x.map(|x| x.to_seconds())),
2431 )
2432 .with_timezone("-05:00"),
2433 ) as ArrayRef
2434 },
2435 ),
2436 ];
2437
2438 resolutions.iter().for_each(|(arrow_type, converter)| {
2439 run_single_column_reader_tests::<Int96Type, _, Int96Type>(
2440 2,
2441 ConvertedType::NONE,
2442 arrow_type.clone(),
2443 converter,
2444 encodings,
2445 );
2446 })
2447 }
2448
2449 #[test]
2450 fn test_int96_from_spark_file_with_provided_schema() {
2451 use arrow_schema::DataType::Timestamp;
2455 let test_data = arrow::util::test_util::parquet_test_data();
2456 let path = format!("{test_data}/int96_from_spark.parquet");
2457 let file = File::open(path).unwrap();
2458
2459 let supplied_schema = Arc::new(Schema::new(vec![Field::new(
2460 "a",
2461 Timestamp(TimeUnit::Microsecond, None),
2462 true,
2463 )]));
2464 let options = ArrowReaderOptions::new().with_schema(supplied_schema.clone());
2465
2466 let mut record_reader =
2467 ParquetRecordBatchReaderBuilder::try_new_with_options(file, options)
2468 .unwrap()
2469 .build()
2470 .unwrap();
2471
2472 let batch = record_reader.next().unwrap().unwrap();
2473 assert_eq!(batch.num_columns(), 1);
2474 let column = batch.column(0);
2475 assert_eq!(column.data_type(), &Timestamp(TimeUnit::Microsecond, None));
2476
2477 let expected = Arc::new(Int64Array::from(vec![
2478 Some(1704141296123456),
2479 Some(1704070800000000),
2480 Some(253402225200000000),
2481 Some(1735599600000000),
2482 None,
2483 Some(9089380393200000000),
2484 ]));
2485
2486 let binding = arrow_cast::cast(batch.column(0), &arrow_schema::DataType::Int64).unwrap();
2491 let casted_timestamps = binding.as_primitive::<types::Int64Type>();
2492
2493 assert_eq!(casted_timestamps.len(), expected.len());
2494
2495 casted_timestamps
2496 .iter()
2497 .zip(expected.iter())
2498 .for_each(|(lhs, rhs)| {
2499 assert_eq!(lhs, rhs);
2500 });
2501 }
2502
2503 #[test]
2504 fn test_int96_from_spark_file_without_provided_schema() {
2505 use arrow_schema::DataType::Timestamp;
2509 let test_data = arrow::util::test_util::parquet_test_data();
2510 let path = format!("{test_data}/int96_from_spark.parquet");
2511 let file = File::open(path).unwrap();
2512
2513 let mut record_reader = ParquetRecordBatchReaderBuilder::try_new(file)
2514 .unwrap()
2515 .build()
2516 .unwrap();
2517
2518 let batch = record_reader.next().unwrap().unwrap();
2519 assert_eq!(batch.num_columns(), 1);
2520 let column = batch.column(0);
2521 assert_eq!(column.data_type(), &Timestamp(TimeUnit::Nanosecond, None));
2522
2523 let expected = Arc::new(Int64Array::from(vec![
2524 Some(1704141296123456000), Some(1704070800000000000), Some(-4852191831933722624), Some(1735599600000000000), None,
2529 Some(-4864435138808946688), ]));
2531
2532 let binding = arrow_cast::cast(batch.column(0), &arrow_schema::DataType::Int64).unwrap();
2537 let casted_timestamps = binding.as_primitive::<types::Int64Type>();
2538
2539 assert_eq!(casted_timestamps.len(), expected.len());
2540
2541 casted_timestamps
2542 .iter()
2543 .zip(expected.iter())
2544 .for_each(|(lhs, rhs)| {
2545 assert_eq!(lhs, rhs);
2546 });
2547 }
2548
2549 struct RandUtf8Gen {}
2550
2551 impl RandGen<ByteArrayType> for RandUtf8Gen {
2552 fn r#gen(len: i32) -> ByteArray {
2553 Int32Type::r#gen(len).to_string().as_str().into()
2554 }
2555 }
2556
2557 #[test]
2558 fn test_utf8_single_column_reader_test() {
2559 fn string_converter<O: OffsetSizeTrait>(vals: &[Option<ByteArray>]) -> ArrayRef {
2560 Arc::new(GenericStringArray::<O>::from_iter(vals.iter().map(|x| {
2561 x.as_ref().map(|b| std::str::from_utf8(b.data()).unwrap())
2562 })))
2563 }
2564
2565 let encodings = &[
2566 Encoding::PLAIN,
2567 Encoding::RLE_DICTIONARY,
2568 Encoding::DELTA_LENGTH_BYTE_ARRAY,
2569 Encoding::DELTA_BYTE_ARRAY,
2570 ];
2571
2572 run_single_column_reader_tests::<ByteArrayType, _, RandUtf8Gen>(
2573 2,
2574 ConvertedType::NONE,
2575 None,
2576 |vals| {
2577 Arc::new(BinaryArray::from_iter(
2578 vals.iter().map(|x| x.as_ref().map(|x| x.data())),
2579 ))
2580 },
2581 encodings,
2582 );
2583
2584 run_single_column_reader_tests::<ByteArrayType, _, RandUtf8Gen>(
2585 2,
2586 ConvertedType::UTF8,
2587 None,
2588 string_converter::<i32>,
2589 encodings,
2590 );
2591
2592 run_single_column_reader_tests::<ByteArrayType, _, RandUtf8Gen>(
2593 2,
2594 ConvertedType::UTF8,
2595 Some(ArrowDataType::Utf8),
2596 string_converter::<i32>,
2597 encodings,
2598 );
2599
2600 run_single_column_reader_tests::<ByteArrayType, _, RandUtf8Gen>(
2601 2,
2602 ConvertedType::UTF8,
2603 Some(ArrowDataType::LargeUtf8),
2604 string_converter::<i64>,
2605 encodings,
2606 );
2607
2608 let small_key_types = [ArrowDataType::Int8, ArrowDataType::UInt8];
2609 for key in &small_key_types {
2610 for encoding in encodings {
2611 let mut opts = TestOptions::new(2, 20, 15).with_null_percent(50);
2612 opts.encoding = *encoding;
2613
2614 let data_type =
2615 ArrowDataType::Dictionary(Box::new(key.clone()), Box::new(ArrowDataType::Utf8));
2616
2617 single_column_reader_test::<ByteArrayType, _, RandUtf8Gen>(
2619 opts,
2620 2,
2621 ConvertedType::UTF8,
2622 Some(data_type.clone()),
2623 move |vals| {
2624 let vals = string_converter::<i32>(vals);
2625 arrow::compute::cast(&vals, &data_type).unwrap()
2626 },
2627 );
2628 }
2629 }
2630
2631 let key_types = [
2632 ArrowDataType::Int16,
2633 ArrowDataType::UInt16,
2634 ArrowDataType::Int32,
2635 ArrowDataType::UInt32,
2636 ArrowDataType::Int64,
2637 ArrowDataType::UInt64,
2638 ];
2639
2640 for key in &key_types {
2641 let data_type =
2642 ArrowDataType::Dictionary(Box::new(key.clone()), Box::new(ArrowDataType::Utf8));
2643
2644 run_single_column_reader_tests::<ByteArrayType, _, RandUtf8Gen>(
2645 2,
2646 ConvertedType::UTF8,
2647 Some(data_type.clone()),
2648 move |vals| {
2649 let vals = string_converter::<i32>(vals);
2650 arrow::compute::cast(&vals, &data_type).unwrap()
2651 },
2652 encodings,
2653 );
2654
2655 let data_type = ArrowDataType::Dictionary(
2656 Box::new(key.clone()),
2657 Box::new(ArrowDataType::LargeUtf8),
2658 );
2659
2660 run_single_column_reader_tests::<ByteArrayType, _, RandUtf8Gen>(
2661 2,
2662 ConvertedType::UTF8,
2663 Some(data_type.clone()),
2664 move |vals| {
2665 let vals = string_converter::<i64>(vals);
2666 arrow::compute::cast(&vals, &data_type).unwrap()
2667 },
2668 encodings,
2669 );
2670 }
2671 }
2672
2673 #[test]
2674 fn test_decimal_nullable_struct() {
2675 let decimals = Decimal256Array::from_iter_values(
2676 [1, 2, 3, 4, 5, 6, 7, 8].into_iter().map(i256::from_i128),
2677 );
2678
2679 let data = ArrayDataBuilder::new(ArrowDataType::Struct(Fields::from(vec![Field::new(
2680 "decimals",
2681 decimals.data_type().clone(),
2682 false,
2683 )])))
2684 .len(8)
2685 .null_bit_buffer(Some(Buffer::from(&[0b11101111])))
2686 .child_data(vec![decimals.into_data()])
2687 .build()
2688 .unwrap();
2689
2690 let written =
2691 RecordBatch::try_from_iter([("struct", Arc::new(StructArray::from(data)) as ArrayRef)])
2692 .unwrap();
2693
2694 let mut buffer = Vec::with_capacity(1024);
2695 let mut writer = ArrowWriter::try_new(&mut buffer, written.schema(), None).unwrap();
2696 writer.write(&written).unwrap();
2697 writer.close().unwrap();
2698
2699 let read = ParquetRecordBatchReader::try_new(Bytes::from(buffer), 3)
2700 .unwrap()
2701 .collect::<Result<Vec<_>, _>>()
2702 .unwrap();
2703
2704 assert_eq!(&written.slice(0, 3), &read[0]);
2705 assert_eq!(&written.slice(3, 3), &read[1]);
2706 assert_eq!(&written.slice(6, 2), &read[2]);
2707 }
2708
2709 #[test]
2710 fn test_int32_nullable_struct() {
2711 let int32 = Int32Array::from_iter_values([1, 2, 3, 4, 5, 6, 7, 8]);
2712 let data = ArrayDataBuilder::new(ArrowDataType::Struct(Fields::from(vec![Field::new(
2713 "int32",
2714 int32.data_type().clone(),
2715 false,
2716 )])))
2717 .len(8)
2718 .null_bit_buffer(Some(Buffer::from(&[0b11101111])))
2719 .child_data(vec![int32.into_data()])
2720 .build()
2721 .unwrap();
2722
2723 let written =
2724 RecordBatch::try_from_iter([("struct", Arc::new(StructArray::from(data)) as ArrayRef)])
2725 .unwrap();
2726
2727 let mut buffer = Vec::with_capacity(1024);
2728 let mut writer = ArrowWriter::try_new(&mut buffer, written.schema(), None).unwrap();
2729 writer.write(&written).unwrap();
2730 writer.close().unwrap();
2731
2732 let read = ParquetRecordBatchReader::try_new(Bytes::from(buffer), 3)
2733 .unwrap()
2734 .collect::<Result<Vec<_>, _>>()
2735 .unwrap();
2736
2737 assert_eq!(&written.slice(0, 3), &read[0]);
2738 assert_eq!(&written.slice(3, 3), &read[1]);
2739 assert_eq!(&written.slice(6, 2), &read[2]);
2740 }
2741
2742 #[test]
2743 fn test_decimal_list() {
2744 let decimals = Decimal128Array::from_iter_values([1, 2, 3, 4, 5, 6, 7, 8]);
2745
2746 let data = ArrayDataBuilder::new(ArrowDataType::List(Arc::new(Field::new_list_field(
2748 decimals.data_type().clone(),
2749 false,
2750 ))))
2751 .len(7)
2752 .add_buffer(Buffer::from_iter([0_i32, 0, 1, 3, 3, 4, 5, 8]))
2753 .null_bit_buffer(Some(Buffer::from(&[0b01010111])))
2754 .child_data(vec![decimals.into_data()])
2755 .build()
2756 .unwrap();
2757
2758 let written =
2759 RecordBatch::try_from_iter([("list", Arc::new(ListArray::from(data)) as ArrayRef)])
2760 .unwrap();
2761
2762 let mut buffer = Vec::with_capacity(1024);
2763 let mut writer = ArrowWriter::try_new(&mut buffer, written.schema(), None).unwrap();
2764 writer.write(&written).unwrap();
2765 writer.close().unwrap();
2766
2767 let read = ParquetRecordBatchReader::try_new(Bytes::from(buffer), 3)
2768 .unwrap()
2769 .collect::<Result<Vec<_>, _>>()
2770 .unwrap();
2771
2772 assert_eq!(&written.slice(0, 3), &read[0]);
2773 assert_eq!(&written.slice(3, 3), &read[1]);
2774 assert_eq!(&written.slice(6, 1), &read[2]);
2775 }
2776
2777 #[test]
2778 fn test_read_decimal_file() {
2779 use arrow_array::Decimal128Array;
2780 let testdata = arrow::util::test_util::parquet_test_data();
2781 let file_variants = vec![
2782 ("byte_array", 4),
2783 ("fixed_length", 25),
2784 ("int32", 4),
2785 ("int64", 10),
2786 ];
2787 for (prefix, target_precision) in file_variants {
2788 let path = format!("{testdata}/{prefix}_decimal.parquet");
2789 let file = File::open(path).unwrap();
2790 let mut record_reader = ParquetRecordBatchReader::try_new(file, 32).unwrap();
2791
2792 let batch = record_reader.next().unwrap().unwrap();
2793 assert_eq!(batch.num_rows(), 24);
2794 let col = batch
2795 .column(0)
2796 .as_any()
2797 .downcast_ref::<Decimal128Array>()
2798 .unwrap();
2799
2800 let expected = 1..25;
2801
2802 assert_eq!(col.precision(), target_precision);
2803 assert_eq!(col.scale(), 2);
2804
2805 for (i, v) in expected.enumerate() {
2806 assert_eq!(col.value(i), v * 100_i128);
2807 }
2808 }
2809 }
2810
2811 #[test]
2812 fn test_read_float16_nonzeros_file() {
2813 use arrow_array::Float16Array;
2814 let testdata = arrow::util::test_util::parquet_test_data();
2815 let path = format!("{testdata}/float16_nonzeros_and_nans.parquet");
2817 let file = File::open(path).unwrap();
2818 let mut record_reader = ParquetRecordBatchReader::try_new(file, 32).unwrap();
2819
2820 let batch = record_reader.next().unwrap().unwrap();
2821 assert_eq!(batch.num_rows(), 8);
2822 let col = batch
2823 .column(0)
2824 .as_any()
2825 .downcast_ref::<Float16Array>()
2826 .unwrap();
2827
2828 let f16_two = f16::ONE + f16::ONE;
2829
2830 assert_eq!(col.null_count(), 1);
2831 assert!(col.is_null(0));
2832 assert_eq!(col.value(1), f16::ONE);
2833 assert_eq!(col.value(2), -f16_two);
2834 assert!(col.value(3).is_nan());
2835 assert_eq!(col.value(4), f16::ZERO);
2836 assert!(col.value(4).is_sign_positive());
2837 assert_eq!(col.value(5), f16::NEG_ONE);
2838 assert_eq!(col.value(6), f16::NEG_ZERO);
2839 assert!(col.value(6).is_sign_negative());
2840 assert_eq!(col.value(7), f16_two);
2841 }
2842
2843 #[test]
2844 fn test_read_float16_zeros_file() {
2845 use arrow_array::Float16Array;
2846 let testdata = arrow::util::test_util::parquet_test_data();
2847 let path = format!("{testdata}/float16_zeros_and_nans.parquet");
2849 let file = File::open(path).unwrap();
2850 let mut record_reader = ParquetRecordBatchReader::try_new(file, 32).unwrap();
2851
2852 let batch = record_reader.next().unwrap().unwrap();
2853 assert_eq!(batch.num_rows(), 3);
2854 let col = batch
2855 .column(0)
2856 .as_any()
2857 .downcast_ref::<Float16Array>()
2858 .unwrap();
2859
2860 assert_eq!(col.null_count(), 1);
2861 assert!(col.is_null(0));
2862 assert_eq!(col.value(1), f16::ZERO);
2863 assert!(col.value(1).is_sign_positive());
2864 assert!(col.value(2).is_nan());
2865 }
2866
2867 #[test]
2868 fn test_read_float32_float64_byte_stream_split() {
2869 let path = format!(
2870 "{}/byte_stream_split.zstd.parquet",
2871 arrow::util::test_util::parquet_test_data(),
2872 );
2873 let file = File::open(path).unwrap();
2874 let record_reader = ParquetRecordBatchReader::try_new(file, 128).unwrap();
2875
2876 let mut row_count = 0;
2877 for batch in record_reader {
2878 let batch = batch.unwrap();
2879 row_count += batch.num_rows();
2880 let f32_col = batch.column(0).as_primitive::<Float32Type>();
2881 let f64_col = batch.column(1).as_primitive::<Float64Type>();
2882
2883 for &x in f32_col.values() {
2885 assert!(x > -10.0);
2886 assert!(x < 10.0);
2887 }
2888 for &x in f64_col.values() {
2889 assert!(x > -10.0);
2890 assert!(x < 10.0);
2891 }
2892 }
2893 assert_eq!(row_count, 300);
2894 }
2895
2896 #[test]
2897 fn test_read_extended_byte_stream_split() {
2898 let path = format!(
2899 "{}/byte_stream_split_extended.gzip.parquet",
2900 arrow::util::test_util::parquet_test_data(),
2901 );
2902 let file = File::open(path).unwrap();
2903 let record_reader = ParquetRecordBatchReader::try_new(file, 128).unwrap();
2904
2905 let mut row_count = 0;
2906 for batch in record_reader {
2907 let batch = batch.unwrap();
2908 row_count += batch.num_rows();
2909
2910 let f16_col = batch.column(0).as_primitive::<Float16Type>();
2912 let f16_bss = batch.column(1).as_primitive::<Float16Type>();
2913 assert_eq!(f16_col.len(), f16_bss.len());
2914 f16_col
2915 .iter()
2916 .zip(f16_bss.iter())
2917 .for_each(|(l, r)| assert_eq!(l.unwrap(), r.unwrap()));
2918
2919 let f32_col = batch.column(2).as_primitive::<Float32Type>();
2921 let f32_bss = batch.column(3).as_primitive::<Float32Type>();
2922 assert_eq!(f32_col.len(), f32_bss.len());
2923 f32_col
2924 .iter()
2925 .zip(f32_bss.iter())
2926 .for_each(|(l, r)| assert_eq!(l.unwrap(), r.unwrap()));
2927
2928 let f64_col = batch.column(4).as_primitive::<Float64Type>();
2930 let f64_bss = batch.column(5).as_primitive::<Float64Type>();
2931 assert_eq!(f64_col.len(), f64_bss.len());
2932 f64_col
2933 .iter()
2934 .zip(f64_bss.iter())
2935 .for_each(|(l, r)| assert_eq!(l.unwrap(), r.unwrap()));
2936
2937 let i32_col = batch.column(6).as_primitive::<types::Int32Type>();
2939 let i32_bss = batch.column(7).as_primitive::<types::Int32Type>();
2940 assert_eq!(i32_col.len(), i32_bss.len());
2941 i32_col
2942 .iter()
2943 .zip(i32_bss.iter())
2944 .for_each(|(l, r)| assert_eq!(l.unwrap(), r.unwrap()));
2945
2946 let i64_col = batch.column(8).as_primitive::<types::Int64Type>();
2948 let i64_bss = batch.column(9).as_primitive::<types::Int64Type>();
2949 assert_eq!(i64_col.len(), i64_bss.len());
2950 i64_col
2951 .iter()
2952 .zip(i64_bss.iter())
2953 .for_each(|(l, r)| assert_eq!(l.unwrap(), r.unwrap()));
2954
2955 let flba_col = batch.column(10).as_fixed_size_binary();
2957 let flba_bss = batch.column(11).as_fixed_size_binary();
2958 assert_eq!(flba_col.len(), flba_bss.len());
2959 flba_col
2960 .iter()
2961 .zip(flba_bss.iter())
2962 .for_each(|(l, r)| assert_eq!(l.unwrap(), r.unwrap()));
2963
2964 let dec_col = batch.column(12).as_primitive::<Decimal128Type>();
2966 let dec_bss = batch.column(13).as_primitive::<Decimal128Type>();
2967 assert_eq!(dec_col.len(), dec_bss.len());
2968 dec_col
2969 .iter()
2970 .zip(dec_bss.iter())
2971 .for_each(|(l, r)| assert_eq!(l.unwrap(), r.unwrap()));
2972 }
2973 assert_eq!(row_count, 200);
2974 }
2975
2976 #[test]
2977 fn test_read_incorrect_map_schema_file() {
2978 let testdata = arrow::util::test_util::parquet_test_data();
2979 let path = format!("{testdata}/incorrect_map_schema.parquet");
2981 let file = File::open(path).unwrap();
2982 let mut record_reader = ParquetRecordBatchReader::try_new(file, 32).unwrap();
2983
2984 let batch = record_reader.next().unwrap().unwrap();
2985 assert_eq!(batch.num_rows(), 1);
2986
2987 let expected_schema = Schema::new(vec![Field::new(
2988 "my_map",
2989 ArrowDataType::Map(
2990 Arc::new(Field::new(
2991 "key_value",
2992 ArrowDataType::Struct(Fields::from(vec![
2993 Field::new("key", ArrowDataType::Utf8, false),
2994 Field::new("value", ArrowDataType::Utf8, true),
2995 ])),
2996 false,
2997 )),
2998 false,
2999 ),
3000 true,
3001 )]);
3002 assert_eq!(batch.schema().as_ref(), &expected_schema);
3003
3004 assert_eq!(batch.num_rows(), 1);
3005 assert_eq!(batch.column(0).null_count(), 0);
3006 assert_eq!(
3007 batch.column(0).as_map().keys().as_ref(),
3008 &StringArray::from(vec!["parent", "name"])
3009 );
3010 assert_eq!(
3011 batch.column(0).as_map().values().as_ref(),
3012 &StringArray::from(vec!["another", "report"])
3013 );
3014 }
3015
3016 #[test]
3017 fn test_read_dict_fixed_size_binary() {
3018 let schema = Arc::new(Schema::new(vec![Field::new(
3019 "a",
3020 ArrowDataType::Dictionary(
3021 Box::new(ArrowDataType::UInt8),
3022 Box::new(ArrowDataType::FixedSizeBinary(8)),
3023 ),
3024 true,
3025 )]));
3026 let keys = UInt8Array::from_iter_values(vec![0, 0, 1]);
3027 let values = FixedSizeBinaryArray::try_from_iter(
3028 vec![
3029 (0u8..8u8).collect::<Vec<u8>>(),
3030 (24u8..32u8).collect::<Vec<u8>>(),
3031 ]
3032 .into_iter(),
3033 )
3034 .unwrap();
3035 let arr = UInt8DictionaryArray::new(keys, Arc::new(values));
3036 let batch = RecordBatch::try_new(schema, vec![Arc::new(arr)]).unwrap();
3037
3038 let mut buffer = Vec::with_capacity(1024);
3039 let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema(), None).unwrap();
3040 writer.write(&batch).unwrap();
3041 writer.close().unwrap();
3042 let read = ParquetRecordBatchReader::try_new(Bytes::from(buffer), 3)
3043 .unwrap()
3044 .collect::<Result<Vec<_>, _>>()
3045 .unwrap();
3046
3047 assert_eq!(read.len(), 1);
3048 assert_eq!(&batch, &read[0])
3049 }
3050
3051 #[test]
3052 fn test_read_nullable_structs_with_binary_dict_as_first_child_column() {
3053 let struct_fields = Fields::from(vec![
3060 Field::new(
3061 "city",
3062 ArrowDataType::Dictionary(
3063 Box::new(ArrowDataType::UInt8),
3064 Box::new(ArrowDataType::Utf8),
3065 ),
3066 true,
3067 ),
3068 Field::new("name", ArrowDataType::Utf8, true),
3069 ]);
3070 let schema = Arc::new(Schema::new(vec![Field::new(
3071 "items",
3072 ArrowDataType::Struct(struct_fields.clone()),
3073 true,
3074 )]));
3075
3076 let items_arr = StructArray::new(
3077 struct_fields,
3078 vec![
3079 Arc::new(DictionaryArray::new(
3080 UInt8Array::from_iter_values(vec![0, 1, 1, 0, 2]),
3081 Arc::new(StringArray::from_iter_values(vec![
3082 "quebec",
3083 "fredericton",
3084 "halifax",
3085 ])),
3086 )),
3087 Arc::new(StringArray::from_iter_values(vec![
3088 "albert", "terry", "lance", "", "tim",
3089 ])),
3090 ],
3091 Some(NullBuffer::from_iter(vec![true, true, true, false, true])),
3092 );
3093
3094 let batch = RecordBatch::try_new(schema, vec![Arc::new(items_arr)]).unwrap();
3095 let mut buffer = Vec::with_capacity(1024);
3096 let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema(), None).unwrap();
3097 writer.write(&batch).unwrap();
3098 writer.close().unwrap();
3099 let read = ParquetRecordBatchReader::try_new(Bytes::from(buffer), 8)
3100 .unwrap()
3101 .collect::<Result<Vec<_>, _>>()
3102 .unwrap();
3103
3104 assert_eq!(read.len(), 1);
3105 assert_eq!(&batch, &read[0])
3106 }
3107
3108 #[derive(Clone)]
3110 struct TestOptions {
3111 num_row_groups: usize,
3114 num_rows: usize,
3116 record_batch_size: usize,
3118 null_percent: Option<usize>,
3120 write_batch_size: usize,
3125 max_data_page_size: usize,
3127 max_dict_page_size: usize,
3129 writer_version: WriterVersion,
3131 enabled_statistics: EnabledStatistics,
3133 encoding: Encoding,
3135 row_selections: Option<(RowSelection, usize)>,
3137 row_filter: Option<Vec<bool>>,
3139 limit: Option<usize>,
3141 offset: Option<usize>,
3143 }
3144
3145 impl std::fmt::Debug for TestOptions {
3147 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
3148 f.debug_struct("TestOptions")
3149 .field("num_row_groups", &self.num_row_groups)
3150 .field("num_rows", &self.num_rows)
3151 .field("record_batch_size", &self.record_batch_size)
3152 .field("null_percent", &self.null_percent)
3153 .field("write_batch_size", &self.write_batch_size)
3154 .field("max_data_page_size", &self.max_data_page_size)
3155 .field("max_dict_page_size", &self.max_dict_page_size)
3156 .field("writer_version", &self.writer_version)
3157 .field("enabled_statistics", &self.enabled_statistics)
3158 .field("encoding", &self.encoding)
3159 .field("row_selections", &self.row_selections.is_some())
3160 .field("row_filter", &self.row_filter.is_some())
3161 .field("limit", &self.limit)
3162 .field("offset", &self.offset)
3163 .finish()
3164 }
3165 }
3166
3167 impl Default for TestOptions {
3168 fn default() -> Self {
3169 Self {
3170 num_row_groups: 2,
3171 num_rows: 100,
3172 record_batch_size: 15,
3173 null_percent: None,
3174 write_batch_size: 64,
3175 max_data_page_size: 1024 * 1024,
3176 max_dict_page_size: 1024 * 1024,
3177 writer_version: WriterVersion::PARQUET_1_0,
3178 enabled_statistics: EnabledStatistics::Page,
3179 encoding: Encoding::PLAIN,
3180 row_selections: None,
3181 row_filter: None,
3182 limit: None,
3183 offset: None,
3184 }
3185 }
3186 }
3187
3188 impl TestOptions {
3189 fn new(num_row_groups: usize, num_rows: usize, record_batch_size: usize) -> Self {
3190 Self {
3191 num_row_groups,
3192 num_rows,
3193 record_batch_size,
3194 ..Default::default()
3195 }
3196 }
3197
3198 fn with_null_percent(self, null_percent: usize) -> Self {
3199 Self {
3200 null_percent: Some(null_percent),
3201 ..self
3202 }
3203 }
3204
3205 fn with_max_data_page_size(self, max_data_page_size: usize) -> Self {
3206 Self {
3207 max_data_page_size,
3208 ..self
3209 }
3210 }
3211
3212 fn with_max_dict_page_size(self, max_dict_page_size: usize) -> Self {
3213 Self {
3214 max_dict_page_size,
3215 ..self
3216 }
3217 }
3218
3219 fn with_enabled_statistics(self, enabled_statistics: EnabledStatistics) -> Self {
3220 Self {
3221 enabled_statistics,
3222 ..self
3223 }
3224 }
3225
3226 fn with_row_selections(self) -> Self {
3227 assert!(self.row_filter.is_none(), "Must set row selection first");
3228
3229 let mut rng = rng();
3230 let step = rng.random_range(self.record_batch_size..self.num_rows);
3231 let row_selections = create_test_selection(
3232 step,
3233 self.num_row_groups * self.num_rows,
3234 rng.random::<bool>(),
3235 );
3236 Self {
3237 row_selections: Some(row_selections),
3238 ..self
3239 }
3240 }
3241
3242 fn with_row_filter(self) -> Self {
3243 let row_count = match &self.row_selections {
3244 Some((_, count)) => *count,
3245 None => self.num_row_groups * self.num_rows,
3246 };
3247
3248 let mut rng = rng();
3249 Self {
3250 row_filter: Some((0..row_count).map(|_| rng.random_bool(0.9)).collect()),
3251 ..self
3252 }
3253 }
3254
3255 fn with_limit(self, limit: usize) -> Self {
3256 Self {
3257 limit: Some(limit),
3258 ..self
3259 }
3260 }
3261
3262 fn with_offset(self, offset: usize) -> Self {
3263 Self {
3264 offset: Some(offset),
3265 ..self
3266 }
3267 }
3268
3269 fn writer_props(&self) -> WriterProperties {
3270 let builder = WriterProperties::builder()
3271 .set_data_page_size_limit(self.max_data_page_size)
3272 .set_write_batch_size(self.write_batch_size)
3273 .set_writer_version(self.writer_version)
3274 .set_statistics_enabled(self.enabled_statistics);
3275
3276 let builder = match self.encoding {
3277 Encoding::RLE_DICTIONARY | Encoding::PLAIN_DICTIONARY => builder
3278 .set_dictionary_enabled(true)
3279 .set_dictionary_page_size_limit(self.max_dict_page_size),
3280 _ => builder
3281 .set_dictionary_enabled(false)
3282 .set_encoding(self.encoding),
3283 };
3284
3285 builder.build()
3286 }
3287 }
3288
3289 fn run_single_column_reader_tests<T, F, G>(
3296 rand_max: i32,
3297 converted_type: ConvertedType,
3298 arrow_type: Option<ArrowDataType>,
3299 converter: F,
3300 encodings: &[Encoding],
3301 ) where
3302 T: DataType,
3303 G: RandGen<T>,
3304 F: Fn(&[Option<T::T>]) -> ArrayRef,
3305 {
3306 let all_options = vec![
3307 TestOptions::new(2, 100, 15),
3310 TestOptions::new(3, 25, 5),
3315 TestOptions::new(4, 100, 25),
3319 TestOptions::new(3, 256, 73).with_max_data_page_size(128),
3321 TestOptions::new(3, 256, 57).with_max_dict_page_size(128),
3323 TestOptions::new(2, 256, 127).with_null_percent(0),
3325 TestOptions::new(2, 256, 93).with_null_percent(25),
3327 TestOptions::new(4, 100, 25).with_limit(0),
3329 TestOptions::new(4, 100, 25).with_limit(50),
3331 TestOptions::new(4, 100, 25).with_limit(10),
3333 TestOptions::new(4, 100, 25).with_limit(101),
3335 TestOptions::new(4, 100, 25).with_offset(30).with_limit(20),
3337 TestOptions::new(4, 100, 25).with_offset(20).with_limit(80),
3339 TestOptions::new(4, 100, 25).with_offset(20).with_limit(81),
3341 TestOptions::new(2, 256, 91)
3343 .with_null_percent(25)
3344 .with_enabled_statistics(EnabledStatistics::Chunk),
3345 TestOptions::new(2, 256, 91)
3347 .with_null_percent(25)
3348 .with_enabled_statistics(EnabledStatistics::None),
3349 TestOptions::new(2, 128, 91)
3351 .with_null_percent(100)
3352 .with_enabled_statistics(EnabledStatistics::None),
3353 TestOptions::new(2, 100, 15).with_row_selections(),
3358 TestOptions::new(3, 25, 5).with_row_selections(),
3363 TestOptions::new(4, 100, 25).with_row_selections(),
3367 TestOptions::new(3, 256, 73)
3369 .with_max_data_page_size(128)
3370 .with_row_selections(),
3371 TestOptions::new(3, 256, 57)
3373 .with_max_dict_page_size(128)
3374 .with_row_selections(),
3375 TestOptions::new(2, 256, 127)
3377 .with_null_percent(0)
3378 .with_row_selections(),
3379 TestOptions::new(2, 256, 93)
3381 .with_null_percent(25)
3382 .with_row_selections(),
3383 TestOptions::new(2, 256, 93)
3385 .with_null_percent(25)
3386 .with_row_selections()
3387 .with_limit(10),
3388 TestOptions::new(2, 256, 93)
3390 .with_null_percent(25)
3391 .with_row_selections()
3392 .with_offset(20)
3393 .with_limit(10),
3394 TestOptions::new(4, 100, 25).with_row_filter(),
3398 TestOptions::new(4, 100, 25)
3400 .with_row_selections()
3401 .with_row_filter(),
3402 TestOptions::new(2, 256, 93)
3404 .with_null_percent(25)
3405 .with_max_data_page_size(10)
3406 .with_row_filter(),
3407 TestOptions::new(2, 256, 93)
3409 .with_null_percent(25)
3410 .with_max_data_page_size(10)
3411 .with_row_selections()
3412 .with_row_filter(),
3413 TestOptions::new(2, 256, 93)
3415 .with_enabled_statistics(EnabledStatistics::None)
3416 .with_max_data_page_size(10)
3417 .with_row_selections(),
3418 ];
3419
3420 all_options.into_iter().for_each(|opts| {
3421 for writer_version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] {
3422 for encoding in encodings {
3423 let opts = TestOptions {
3424 writer_version,
3425 encoding: *encoding,
3426 ..opts.clone()
3427 };
3428
3429 single_column_reader_test::<T, _, G>(
3430 opts,
3431 rand_max,
3432 converted_type,
3433 arrow_type.clone(),
3434 &converter,
3435 )
3436 }
3437 }
3438 });
3439 }
3440
3441 fn single_column_reader_test<T, F, G>(
3445 opts: TestOptions,
3446 rand_max: i32,
3447 converted_type: ConvertedType,
3448 arrow_type: Option<ArrowDataType>,
3449 converter: F,
3450 ) where
3451 T: DataType,
3452 G: RandGen<T>,
3453 F: Fn(&[Option<T::T>]) -> ArrayRef,
3454 {
3455 println!(
3457 "Running type {:?} single_column_reader_test ConvertedType::{}/ArrowType::{:?} with Options: {:?}",
3458 T::get_physical_type(),
3459 converted_type,
3460 arrow_type,
3461 opts
3462 );
3463
3464 let (repetition, def_levels) = match opts.null_percent.as_ref() {
3466 Some(null_percent) => {
3467 let mut rng = rng();
3468
3469 let def_levels: Vec<Vec<i16>> = (0..opts.num_row_groups)
3470 .map(|_| {
3471 std::iter::from_fn(|| {
3472 Some((rng.next_u32() as usize % 100 >= *null_percent) as i16)
3473 })
3474 .take(opts.num_rows)
3475 .collect()
3476 })
3477 .collect();
3478 (Repetition::OPTIONAL, Some(def_levels))
3479 }
3480 None => (Repetition::REQUIRED, None),
3481 };
3482
3483 let values: Vec<Vec<T::T>> = (0..opts.num_row_groups)
3485 .map(|idx| {
3486 let null_count = match def_levels.as_ref() {
3487 Some(d) => d[idx].iter().filter(|x| **x == 0).count(),
3488 None => 0,
3489 };
3490 G::gen_vec(rand_max, opts.num_rows - null_count)
3491 })
3492 .collect();
3493
3494 let len = match T::get_physical_type() {
3495 crate::basic::Type::FIXED_LEN_BYTE_ARRAY => rand_max,
3496 crate::basic::Type::INT96 => 12,
3497 _ => -1,
3498 };
3499
3500 let fields = vec![Arc::new(
3501 Type::primitive_type_builder("leaf", T::get_physical_type())
3502 .with_repetition(repetition)
3503 .with_converted_type(converted_type)
3504 .with_length(len)
3505 .build()
3506 .unwrap(),
3507 )];
3508
3509 let schema = Arc::new(
3510 Type::group_type_builder("test_schema")
3511 .with_fields(fields)
3512 .build()
3513 .unwrap(),
3514 );
3515
3516 let arrow_field = arrow_type.map(|t| Field::new("leaf", t, false));
3517
3518 let mut file = tempfile::tempfile().unwrap();
3519
3520 generate_single_column_file_with_data::<T>(
3521 &values,
3522 def_levels.as_ref(),
3523 file.try_clone().unwrap(), schema,
3525 arrow_field,
3526 &opts,
3527 )
3528 .unwrap();
3529
3530 file.rewind().unwrap();
3531
3532 let options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::from(
3533 opts.enabled_statistics == EnabledStatistics::Page,
3534 ));
3535
3536 let mut builder =
3537 ParquetRecordBatchReaderBuilder::try_new_with_options(file, options).unwrap();
3538
3539 let expected_data = match opts.row_selections {
3540 Some((selections, row_count)) => {
3541 let mut without_skip_data = gen_expected_data::<T>(def_levels.as_ref(), &values);
3542
3543 let mut skip_data: Vec<Option<T::T>> = vec![];
3544 let dequeue: VecDeque<RowSelector> = selections.clone().into();
3545 for select in dequeue {
3546 if select.skip {
3547 without_skip_data.drain(0..select.row_count);
3548 } else {
3549 skip_data.extend(without_skip_data.drain(0..select.row_count));
3550 }
3551 }
3552 builder = builder.with_row_selection(selections);
3553
3554 assert_eq!(skip_data.len(), row_count);
3555 skip_data
3556 }
3557 None => {
3558 let expected_data = gen_expected_data::<T>(def_levels.as_ref(), &values);
3560 assert_eq!(expected_data.len(), opts.num_rows * opts.num_row_groups);
3561 expected_data
3562 }
3563 };
3564
3565 let mut expected_data = match opts.row_filter {
3566 Some(filter) => {
3567 let expected_data = expected_data
3568 .into_iter()
3569 .zip(filter.iter())
3570 .filter_map(|(d, f)| f.then(|| d))
3571 .collect();
3572
3573 let mut filter_offset = 0;
3574 let filter = RowFilter::new(vec![Box::new(ArrowPredicateFn::new(
3575 ProjectionMask::all(),
3576 move |b| {
3577 let array = BooleanArray::from_iter(
3578 filter
3579 .iter()
3580 .skip(filter_offset)
3581 .take(b.num_rows())
3582 .map(|x| Some(*x)),
3583 );
3584 filter_offset += b.num_rows();
3585 Ok(array)
3586 },
3587 ))]);
3588
3589 builder = builder.with_row_filter(filter);
3590 expected_data
3591 }
3592 None => expected_data,
3593 };
3594
3595 if let Some(offset) = opts.offset {
3596 builder = builder.with_offset(offset);
3597 expected_data = expected_data.into_iter().skip(offset).collect();
3598 }
3599
3600 if let Some(limit) = opts.limit {
3601 builder = builder.with_limit(limit);
3602 expected_data = expected_data.into_iter().take(limit).collect();
3603 }
3604
3605 let mut record_reader = builder
3606 .with_batch_size(opts.record_batch_size)
3607 .build()
3608 .unwrap();
3609
3610 let mut total_read = 0;
3611 loop {
3612 let maybe_batch = record_reader.next();
3613 if total_read < expected_data.len() {
3614 let end = min(total_read + opts.record_batch_size, expected_data.len());
3615 let batch = maybe_batch.unwrap().unwrap();
3616 assert_eq!(end - total_read, batch.num_rows());
3617
3618 let a = converter(&expected_data[total_read..end]);
3619 let b = batch.column(0);
3620
3621 assert_eq!(a.data_type(), b.data_type());
3622 assert_eq!(a.to_data(), b.to_data());
3623 assert_eq!(
3624 a.as_any().type_id(),
3625 b.as_any().type_id(),
3626 "incorrect type ids"
3627 );
3628
3629 total_read = end;
3630 } else {
3631 assert!(maybe_batch.is_none());
3632 break;
3633 }
3634 }
3635 }
3636
3637 fn gen_expected_data<T: DataType>(
3638 def_levels: Option<&Vec<Vec<i16>>>,
3639 values: &[Vec<T::T>],
3640 ) -> Vec<Option<T::T>> {
3641 let data: Vec<Option<T::T>> = match def_levels {
3642 Some(levels) => {
3643 let mut values_iter = values.iter().flatten();
3644 levels
3645 .iter()
3646 .flatten()
3647 .map(|d| match d {
3648 1 => Some(values_iter.next().cloned().unwrap()),
3649 0 => None,
3650 _ => unreachable!(),
3651 })
3652 .collect()
3653 }
3654 None => values.iter().flatten().cloned().map(Some).collect(),
3655 };
3656 data
3657 }
3658
3659 fn generate_single_column_file_with_data<T: DataType>(
3660 values: &[Vec<T::T>],
3661 def_levels: Option<&Vec<Vec<i16>>>,
3662 file: File,
3663 schema: TypePtr,
3664 field: Option<Field>,
3665 opts: &TestOptions,
3666 ) -> Result<ParquetMetaData> {
3667 let mut writer_props = opts.writer_props();
3668 if let Some(field) = field {
3669 let arrow_schema = Schema::new(vec![field]);
3670 add_encoded_arrow_schema_to_metadata(&arrow_schema, &mut writer_props);
3671 }
3672
3673 let mut writer = SerializedFileWriter::new(file, schema, Arc::new(writer_props))?;
3674
3675 for (idx, v) in values.iter().enumerate() {
3676 let def_levels = def_levels.map(|d| d[idx].as_slice());
3677 let mut row_group_writer = writer.next_row_group()?;
3678 {
3679 let mut column_writer = row_group_writer
3680 .next_column()?
3681 .expect("Column writer is none!");
3682
3683 column_writer
3684 .typed::<T>()
3685 .write_batch(v, def_levels, None)?;
3686
3687 column_writer.close()?;
3688 }
3689 row_group_writer.close()?;
3690 }
3691
3692 writer.close()
3693 }
3694
3695 fn get_test_file(file_name: &str) -> File {
3696 let mut path = PathBuf::new();
3697 path.push(arrow::util::test_util::arrow_test_data());
3698 path.push(file_name);
3699
3700 File::open(path.as_path()).expect("File not found!")
3701 }
3702
3703 #[test]
3704 fn test_read_structs() {
3705 let testdata = arrow::util::test_util::parquet_test_data();
3709 let path = format!("{testdata}/nested_structs.rust.parquet");
3710 let file = File::open(&path).unwrap();
3711 let record_batch_reader = ParquetRecordBatchReader::try_new(file, 60).unwrap();
3712
3713 for batch in record_batch_reader {
3714 batch.unwrap();
3715 }
3716
3717 let file = File::open(&path).unwrap();
3718 let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
3719
3720 let mask = ProjectionMask::leaves(builder.parquet_schema(), [3, 8, 10]);
3721 let projected_reader = builder
3722 .with_projection(mask)
3723 .with_batch_size(60)
3724 .build()
3725 .unwrap();
3726
3727 let expected_schema = Schema::new(vec![
3728 Field::new(
3729 "roll_num",
3730 ArrowDataType::Struct(Fields::from(vec![Field::new(
3731 "count",
3732 ArrowDataType::UInt64,
3733 false,
3734 )])),
3735 false,
3736 ),
3737 Field::new(
3738 "PC_CUR",
3739 ArrowDataType::Struct(Fields::from(vec![
3740 Field::new("mean", ArrowDataType::Int64, false),
3741 Field::new("sum", ArrowDataType::Int64, false),
3742 ])),
3743 false,
3744 ),
3745 ]);
3746
3747 assert_eq!(&expected_schema, projected_reader.schema().as_ref());
3749
3750 for batch in projected_reader {
3751 let batch = batch.unwrap();
3752 assert_eq!(batch.schema().as_ref(), &expected_schema);
3753 }
3754 }
3755
3756 #[test]
3757 fn test_read_structs_by_name() {
3759 let testdata = arrow::util::test_util::parquet_test_data();
3760 let path = format!("{testdata}/nested_structs.rust.parquet");
3761 let file = File::open(&path).unwrap();
3762 let record_batch_reader = ParquetRecordBatchReader::try_new(file, 60).unwrap();
3763
3764 for batch in record_batch_reader {
3765 batch.unwrap();
3766 }
3767
3768 let file = File::open(&path).unwrap();
3769 let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
3770
3771 let mask = ProjectionMask::columns(
3772 builder.parquet_schema(),
3773 ["roll_num.count", "PC_CUR.mean", "PC_CUR.sum"],
3774 );
3775 let projected_reader = builder
3776 .with_projection(mask)
3777 .with_batch_size(60)
3778 .build()
3779 .unwrap();
3780
3781 let expected_schema = Schema::new(vec![
3782 Field::new(
3783 "roll_num",
3784 ArrowDataType::Struct(Fields::from(vec![Field::new(
3785 "count",
3786 ArrowDataType::UInt64,
3787 false,
3788 )])),
3789 false,
3790 ),
3791 Field::new(
3792 "PC_CUR",
3793 ArrowDataType::Struct(Fields::from(vec![
3794 Field::new("mean", ArrowDataType::Int64, false),
3795 Field::new("sum", ArrowDataType::Int64, false),
3796 ])),
3797 false,
3798 ),
3799 ]);
3800
3801 assert_eq!(&expected_schema, projected_reader.schema().as_ref());
3802
3803 for batch in projected_reader {
3804 let batch = batch.unwrap();
3805 assert_eq!(batch.schema().as_ref(), &expected_schema);
3806 }
3807 }
3808
3809 #[test]
3810 fn test_read_maps() {
3811 let testdata = arrow::util::test_util::parquet_test_data();
3812 let path = format!("{testdata}/nested_maps.snappy.parquet");
3813 let file = File::open(path).unwrap();
3814 let record_batch_reader = ParquetRecordBatchReader::try_new(file, 60).unwrap();
3815
3816 for batch in record_batch_reader {
3817 batch.unwrap();
3818 }
3819 }
3820
3821 #[test]
3823 fn test_unknown_logical_type() {
3824 let message_type = "message uk {
3825 OPTIONAL INT32 uki32 (UNKNOWN);
3826 OPTIONAL INT64 uki64 (UNKNOWN);
3827 OPTIONAL INT96 uki96 (UNKNOWN);
3828 OPTIONAL BOOLEAN ukbool (UNKNOWN);
3829 OPTIONAL FLOAT ukfloat (UNKNOWN);
3830 OPTIONAL DOUBLE ukdbl (UNKNOWN);
3831 OPTIONAL BYTE_ARRAY ukbytes (UNKNOWN);
3832 OPTIONAL FIXED_LEN_BYTE_ARRAY(10) ukflba (UNKNOWN);
3833 }";
3834
3835 let schema = Arc::new(parse_message_type(message_type).unwrap());
3836 let file = tempfile::tempfile().unwrap();
3837
3838 let mut writer =
3839 SerializedFileWriter::new(file.try_clone().unwrap(), schema, Default::default())
3840 .unwrap();
3841
3842 let mut row_group_writer = writer.next_row_group().unwrap();
3843
3844 fn write_nulls<T: DataType>(row_group_writer: &mut SerializedRowGroupWriter<'_, File>) {
3845 let mut column_writer = row_group_writer.next_column().unwrap().unwrap();
3846 column_writer
3848 .typed::<T>()
3849 .write_batch(&[], Some(&[0, 0, 0, 0]), None)
3850 .unwrap();
3851 column_writer.close().unwrap();
3852 }
3853
3854 write_nulls::<Int32Type>(&mut row_group_writer);
3856
3857 write_nulls::<Int64Type>(&mut row_group_writer);
3859
3860 write_nulls::<Int96Type>(&mut row_group_writer);
3862
3863 write_nulls::<BoolType>(&mut row_group_writer);
3865
3866 write_nulls::<FloatType>(&mut row_group_writer);
3868
3869 write_nulls::<DoubleType>(&mut row_group_writer);
3871
3872 write_nulls::<ByteArrayType>(&mut row_group_writer);
3874
3875 write_nulls::<FixedLenByteArrayType>(&mut row_group_writer);
3877
3878 row_group_writer.close().unwrap();
3879
3880 writer.close().unwrap();
3881
3882 let mut reader = ParquetRecordBatchReader::try_new(file, 4).unwrap();
3883 let batch = reader.next().unwrap().unwrap();
3884
3885 for col in batch.columns() {
3886 assert_eq!(col.len(), 4);
3887 assert_eq!(col.logical_null_count(), 4);
3888 assert_eq!(*col.data_type(), ArrowDataType::Null);
3889 }
3890 }
3891
3892 #[test]
3893 fn test_nested_nullability() {
3894 let message_type = "message nested {
3895 OPTIONAL Group group {
3896 REQUIRED INT32 leaf;
3897 }
3898 }";
3899
3900 let file = tempfile::tempfile().unwrap();
3901 let schema = Arc::new(parse_message_type(message_type).unwrap());
3902
3903 {
3904 let mut writer =
3906 SerializedFileWriter::new(file.try_clone().unwrap(), schema, Default::default())
3907 .unwrap();
3908
3909 {
3910 let mut row_group_writer = writer.next_row_group().unwrap();
3911 let mut column_writer = row_group_writer.next_column().unwrap().unwrap();
3912
3913 column_writer
3914 .typed::<Int32Type>()
3915 .write_batch(&[34, 76], Some(&[0, 1, 0, 1]), None)
3916 .unwrap();
3917
3918 column_writer.close().unwrap();
3919 row_group_writer.close().unwrap();
3920 }
3921
3922 writer.close().unwrap();
3923 }
3924
3925 let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
3926 let mask = ProjectionMask::leaves(builder.parquet_schema(), [0]);
3927
3928 let reader = builder.with_projection(mask).build().unwrap();
3929
3930 let expected_schema = Schema::new(vec![Field::new(
3931 "group",
3932 ArrowDataType::Struct(vec![Field::new("leaf", ArrowDataType::Int32, false)].into()),
3933 true,
3934 )]);
3935
3936 let batch = reader.into_iter().next().unwrap().unwrap();
3937 assert_eq!(batch.schema().as_ref(), &expected_schema);
3938 assert_eq!(batch.num_rows(), 4);
3939 assert_eq!(batch.column(0).null_count(), 2);
3940 }
3941
3942 #[test]
3943 fn test_dictionary_preservation() {
3944 let fields = vec![Arc::new(
3945 Type::primitive_type_builder("leaf", PhysicalType::BYTE_ARRAY)
3946 .with_repetition(Repetition::OPTIONAL)
3947 .with_converted_type(ConvertedType::UTF8)
3948 .build()
3949 .unwrap(),
3950 )];
3951
3952 let schema = Arc::new(
3953 Type::group_type_builder("test_schema")
3954 .with_fields(fields)
3955 .build()
3956 .unwrap(),
3957 );
3958
3959 let dict_type = ArrowDataType::Dictionary(
3960 Box::new(ArrowDataType::Int32),
3961 Box::new(ArrowDataType::Utf8),
3962 );
3963
3964 let arrow_field = Field::new("leaf", dict_type, true);
3965
3966 let mut file = tempfile::tempfile().unwrap();
3967
3968 let values = vec![
3969 vec![
3970 ByteArray::from("hello"),
3971 ByteArray::from("a"),
3972 ByteArray::from("b"),
3973 ByteArray::from("d"),
3974 ],
3975 vec![
3976 ByteArray::from("c"),
3977 ByteArray::from("a"),
3978 ByteArray::from("b"),
3979 ],
3980 ];
3981
3982 let def_levels = vec![
3983 vec![1, 0, 0, 1, 0, 0, 1, 1],
3984 vec![0, 0, 1, 1, 0, 0, 1, 0, 0],
3985 ];
3986
3987 let opts = TestOptions {
3988 encoding: Encoding::RLE_DICTIONARY,
3989 ..Default::default()
3990 };
3991
3992 generate_single_column_file_with_data::<ByteArrayType>(
3993 &values,
3994 Some(&def_levels),
3995 file.try_clone().unwrap(), schema,
3997 Some(arrow_field),
3998 &opts,
3999 )
4000 .unwrap();
4001
4002 file.rewind().unwrap();
4003
4004 let record_reader = ParquetRecordBatchReader::try_new(file, 3).unwrap();
4005
4006 let batches = record_reader
4007 .collect::<Result<Vec<RecordBatch>, _>>()
4008 .unwrap();
4009
4010 assert_eq!(batches.len(), 6);
4011 assert!(batches.iter().all(|x| x.num_columns() == 1));
4012
4013 let row_counts = batches
4014 .iter()
4015 .map(|x| (x.num_rows(), x.column(0).null_count()))
4016 .collect::<Vec<_>>();
4017
4018 assert_eq!(
4019 row_counts,
4020 vec![(3, 2), (3, 2), (3, 1), (3, 1), (3, 2), (2, 2)]
4021 );
4022
4023 let get_dict = |batch: &RecordBatch| batch.column(0).to_data().child_data()[0].clone();
4024
4025 assert_eq!(get_dict(&batches[0]), get_dict(&batches[1]));
4027 assert_ne!(get_dict(&batches[1]), get_dict(&batches[2]));
4029 assert_ne!(get_dict(&batches[2]), get_dict(&batches[3]));
4030 assert_eq!(get_dict(&batches[3]), get_dict(&batches[4]));
4032 assert_eq!(get_dict(&batches[4]), get_dict(&batches[5]));
4033 }
4034
4035 #[test]
4036 fn test_read_null_list() {
4037 let testdata = arrow::util::test_util::parquet_test_data();
4038 let path = format!("{testdata}/null_list.parquet");
4039 let file = File::open(path).unwrap();
4040 let mut record_batch_reader = ParquetRecordBatchReader::try_new(file, 60).unwrap();
4041
4042 let batch = record_batch_reader.next().unwrap().unwrap();
4043 assert_eq!(batch.num_rows(), 1);
4044 assert_eq!(batch.num_columns(), 1);
4045 assert_eq!(batch.column(0).len(), 1);
4046
4047 let list = batch
4048 .column(0)
4049 .as_any()
4050 .downcast_ref::<ListArray>()
4051 .unwrap();
4052 assert_eq!(list.len(), 1);
4053 assert!(list.is_valid(0));
4054
4055 let val = list.value(0);
4056 assert_eq!(val.len(), 0);
4057 }
4058
4059 #[test]
4060 fn test_null_schema_inference() {
4061 let testdata = arrow::util::test_util::parquet_test_data();
4062 let path = format!("{testdata}/null_list.parquet");
4063 let file = File::open(path).unwrap();
4064
4065 let arrow_field = Field::new(
4066 "emptylist",
4067 ArrowDataType::List(Arc::new(Field::new_list_field(ArrowDataType::Null, true))),
4068 true,
4069 );
4070
4071 let options = ArrowReaderOptions::new().with_skip_arrow_metadata(true);
4072 let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(file, options).unwrap();
4073 let schema = builder.schema();
4074 assert_eq!(schema.fields().len(), 1);
4075 assert_eq!(schema.field(0), &arrow_field);
4076 }
4077
4078 #[test]
4079 fn test_skip_metadata() {
4080 let col = Arc::new(TimestampNanosecondArray::from_iter_values(vec![0, 1, 2]));
4081 let field = Field::new("col", col.data_type().clone(), true);
4082
4083 let schema_without_metadata = Arc::new(Schema::new(vec![field.clone()]));
4084
4085 let metadata = [("key".to_string(), "value".to_string())]
4086 .into_iter()
4087 .collect();
4088
4089 let schema_with_metadata = Arc::new(Schema::new(vec![field.with_metadata(metadata)]));
4090
4091 assert_ne!(schema_with_metadata, schema_without_metadata);
4092
4093 let batch =
4094 RecordBatch::try_new(schema_with_metadata.clone(), vec![col as ArrayRef]).unwrap();
4095
4096 let file = |version: WriterVersion| {
4097 let props = WriterProperties::builder()
4098 .set_writer_version(version)
4099 .build();
4100
4101 let file = tempfile().unwrap();
4102 let mut writer =
4103 ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props))
4104 .unwrap();
4105 writer.write(&batch).unwrap();
4106 writer.close().unwrap();
4107 file
4108 };
4109
4110 let skip_options = ArrowReaderOptions::new().with_skip_arrow_metadata(true);
4111
4112 let v1_reader = file(WriterVersion::PARQUET_1_0);
4113 let v2_reader = file(WriterVersion::PARQUET_2_0);
4114
4115 let arrow_reader =
4116 ParquetRecordBatchReader::try_new(v1_reader.try_clone().unwrap(), 1024).unwrap();
4117 assert_eq!(arrow_reader.schema(), schema_with_metadata);
4118
4119 let reader =
4120 ParquetRecordBatchReaderBuilder::try_new_with_options(v1_reader, skip_options.clone())
4121 .unwrap()
4122 .build()
4123 .unwrap();
4124 assert_eq!(reader.schema(), schema_without_metadata);
4125
4126 let arrow_reader =
4127 ParquetRecordBatchReader::try_new(v2_reader.try_clone().unwrap(), 1024).unwrap();
4128 assert_eq!(arrow_reader.schema(), schema_with_metadata);
4129
4130 let reader = ParquetRecordBatchReaderBuilder::try_new_with_options(v2_reader, skip_options)
4131 .unwrap()
4132 .build()
4133 .unwrap();
4134 assert_eq!(reader.schema(), schema_without_metadata);
4135 }
4136
4137 fn write_parquet_from_iter<I, F>(value: I) -> File
4138 where
4139 I: IntoIterator<Item = (F, ArrayRef)>,
4140 F: AsRef<str>,
4141 {
4142 let batch = RecordBatch::try_from_iter(value).unwrap();
4143 let file = tempfile().unwrap();
4144 let mut writer =
4145 ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema().clone(), None).unwrap();
4146 writer.write(&batch).unwrap();
4147 writer.close().unwrap();
4148 file
4149 }
4150
4151 fn run_schema_test_with_error<I, F>(value: I, schema: SchemaRef, expected_error: &str)
4152 where
4153 I: IntoIterator<Item = (F, ArrayRef)>,
4154 F: AsRef<str>,
4155 {
4156 let file = write_parquet_from_iter(value);
4157 let options_with_schema = ArrowReaderOptions::new().with_schema(schema.clone());
4158 let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(
4159 file.try_clone().unwrap(),
4160 options_with_schema,
4161 );
4162 assert_eq!(builder.err().unwrap().to_string(), expected_error);
4163 }
4164
4165 #[test]
4166 fn test_schema_too_few_columns() {
4167 run_schema_test_with_error(
4168 vec![
4169 ("int64", Arc::new(Int64Array::from(vec![0])) as ArrayRef),
4170 ("int32", Arc::new(Int32Array::from(vec![0])) as ArrayRef),
4171 ],
4172 Arc::new(Schema::new(vec![Field::new(
4173 "int64",
4174 ArrowDataType::Int64,
4175 false,
4176 )])),
4177 "Arrow: incompatible arrow schema, expected 2 struct fields got 1",
4178 );
4179 }
4180
4181 #[test]
4182 fn test_schema_too_many_columns() {
4183 run_schema_test_with_error(
4184 vec![("int64", Arc::new(Int64Array::from(vec![0])) as ArrayRef)],
4185 Arc::new(Schema::new(vec![
4186 Field::new("int64", ArrowDataType::Int64, false),
4187 Field::new("int32", ArrowDataType::Int32, false),
4188 ])),
4189 "Arrow: incompatible arrow schema, expected 1 struct fields got 2",
4190 );
4191 }
4192
4193 #[test]
4194 fn test_schema_mismatched_column_names() {
4195 run_schema_test_with_error(
4196 vec![("int64", Arc::new(Int64Array::from(vec![0])) as ArrayRef)],
4197 Arc::new(Schema::new(vec![Field::new(
4198 "other",
4199 ArrowDataType::Int64,
4200 false,
4201 )])),
4202 "Arrow: incompatible arrow schema, expected field named int64 got other",
4203 );
4204 }
4205
4206 #[test]
4207 fn test_schema_incompatible_columns() {
4208 run_schema_test_with_error(
4209 vec![
4210 (
4211 "col1_invalid",
4212 Arc::new(Int64Array::from(vec![0])) as ArrayRef,
4213 ),
4214 (
4215 "col2_valid",
4216 Arc::new(Int32Array::from(vec![0])) as ArrayRef,
4217 ),
4218 (
4219 "col3_invalid",
4220 Arc::new(Date64Array::from(vec![0])) as ArrayRef,
4221 ),
4222 ],
4223 Arc::new(Schema::new(vec![
4224 Field::new("col1_invalid", ArrowDataType::Int32, false),
4225 Field::new("col2_valid", ArrowDataType::Int32, false),
4226 Field::new("col3_invalid", ArrowDataType::Int32, false),
4227 ])),
4228 "Arrow: Incompatible supplied Arrow schema: data type mismatch for field col1_invalid: requested Int32 but found Int64, data type mismatch for field col3_invalid: requested Int32 but found Int64",
4229 );
4230 }
4231
4232 #[test]
4233 fn test_one_incompatible_nested_column() {
4234 let nested_fields = Fields::from(vec![
4235 Field::new("nested1_valid", ArrowDataType::Utf8, false),
4236 Field::new("nested1_invalid", ArrowDataType::Int64, false),
4237 ]);
4238 let nested = StructArray::try_new(
4239 nested_fields,
4240 vec![
4241 Arc::new(StringArray::from(vec!["a"])) as ArrayRef,
4242 Arc::new(Int64Array::from(vec![0])) as ArrayRef,
4243 ],
4244 None,
4245 )
4246 .expect("struct array");
4247 let supplied_nested_fields = Fields::from(vec![
4248 Field::new("nested1_valid", ArrowDataType::Utf8, false),
4249 Field::new("nested1_invalid", ArrowDataType::Int32, false),
4250 ]);
4251 run_schema_test_with_error(
4252 vec![
4253 ("col1", Arc::new(Int64Array::from(vec![0])) as ArrayRef),
4254 ("col2", Arc::new(Int32Array::from(vec![0])) as ArrayRef),
4255 ("nested", Arc::new(nested) as ArrayRef),
4256 ],
4257 Arc::new(Schema::new(vec![
4258 Field::new("col1", ArrowDataType::Int64, false),
4259 Field::new("col2", ArrowDataType::Int32, false),
4260 Field::new(
4261 "nested",
4262 ArrowDataType::Struct(supplied_nested_fields),
4263 false,
4264 ),
4265 ])),
4266 "Arrow: Incompatible supplied Arrow schema: data type mismatch for field nested: \
4267 requested Struct(\"nested1_valid\": non-null Utf8, \"nested1_invalid\": non-null Int32) \
4268 but found Struct(\"nested1_valid\": non-null Utf8, \"nested1_invalid\": non-null Int64)",
4269 );
4270 }
4271
4272 fn utf8_parquet() -> Bytes {
4274 let input = StringArray::from_iter_values(vec!["foo", "bar", "baz"]);
4275 let batch = RecordBatch::try_from_iter(vec![("column1", Arc::new(input) as _)]).unwrap();
4276 let props = None;
4277 let mut parquet_data = vec![];
4279 let mut writer = ArrowWriter::try_new(&mut parquet_data, batch.schema(), props).unwrap();
4280 writer.write(&batch).unwrap();
4281 writer.close().unwrap();
4282 Bytes::from(parquet_data)
4283 }
4284
4285 #[test]
4286 fn test_schema_error_bad_types() {
4287 let parquet_data = utf8_parquet();
4289
4290 let input_schema: SchemaRef = Arc::new(Schema::new(vec![Field::new(
4292 "column1",
4293 arrow::datatypes::DataType::Int32,
4294 false,
4295 )]));
4296
4297 let reader_options = ArrowReaderOptions::new().with_schema(input_schema.clone());
4299 let err =
4300 ParquetRecordBatchReaderBuilder::try_new_with_options(parquet_data, reader_options)
4301 .unwrap_err();
4302 assert_eq!(
4303 err.to_string(),
4304 "Arrow: Incompatible supplied Arrow schema: data type mismatch for field column1: requested Int32 but found Utf8"
4305 )
4306 }
4307
4308 #[test]
4309 fn test_schema_error_bad_nullability() {
4310 let parquet_data = utf8_parquet();
4312
4313 let input_schema: SchemaRef = Arc::new(Schema::new(vec![Field::new(
4315 "column1",
4316 arrow::datatypes::DataType::Utf8,
4317 true,
4318 )]));
4319
4320 let reader_options = ArrowReaderOptions::new().with_schema(input_schema.clone());
4322 let err =
4323 ParquetRecordBatchReaderBuilder::try_new_with_options(parquet_data, reader_options)
4324 .unwrap_err();
4325 assert_eq!(
4326 err.to_string(),
4327 "Arrow: Incompatible supplied Arrow schema: nullability mismatch for field column1: expected true but found false"
4328 )
4329 }
4330
4331 #[test]
4332 fn test_read_binary_as_utf8() {
4333 let file = write_parquet_from_iter(vec![
4334 (
4335 "binary_to_utf8",
4336 Arc::new(BinaryArray::from(vec![
4337 b"one".as_ref(),
4338 b"two".as_ref(),
4339 b"three".as_ref(),
4340 ])) as ArrayRef,
4341 ),
4342 (
4343 "large_binary_to_large_utf8",
4344 Arc::new(LargeBinaryArray::from(vec![
4345 b"one".as_ref(),
4346 b"two".as_ref(),
4347 b"three".as_ref(),
4348 ])) as ArrayRef,
4349 ),
4350 (
4351 "binary_view_to_utf8_view",
4352 Arc::new(BinaryViewArray::from(vec![
4353 b"one".as_ref(),
4354 b"two".as_ref(),
4355 b"three".as_ref(),
4356 ])) as ArrayRef,
4357 ),
4358 ]);
4359 let supplied_fields = Fields::from(vec![
4360 Field::new("binary_to_utf8", ArrowDataType::Utf8, false),
4361 Field::new(
4362 "large_binary_to_large_utf8",
4363 ArrowDataType::LargeUtf8,
4364 false,
4365 ),
4366 Field::new("binary_view_to_utf8_view", ArrowDataType::Utf8View, false),
4367 ]);
4368
4369 let options = ArrowReaderOptions::new().with_schema(Arc::new(Schema::new(supplied_fields)));
4370 let mut arrow_reader = ParquetRecordBatchReaderBuilder::try_new_with_options(
4371 file.try_clone().unwrap(),
4372 options,
4373 )
4374 .expect("reader builder with schema")
4375 .build()
4376 .expect("reader with schema");
4377
4378 let batch = arrow_reader.next().unwrap().unwrap();
4379 assert_eq!(batch.num_columns(), 3);
4380 assert_eq!(batch.num_rows(), 3);
4381 assert_eq!(
4382 batch
4383 .column(0)
4384 .as_string::<i32>()
4385 .iter()
4386 .collect::<Vec<_>>(),
4387 vec![Some("one"), Some("two"), Some("three")]
4388 );
4389
4390 assert_eq!(
4391 batch
4392 .column(1)
4393 .as_string::<i64>()
4394 .iter()
4395 .collect::<Vec<_>>(),
4396 vec![Some("one"), Some("two"), Some("three")]
4397 );
4398
4399 assert_eq!(
4400 batch.column(2).as_string_view().iter().collect::<Vec<_>>(),
4401 vec![Some("one"), Some("two"), Some("three")]
4402 );
4403 }
4404
4405 #[test]
4406 #[should_panic(expected = "Invalid UTF8 sequence at")]
4407 fn test_read_non_utf8_binary_as_utf8() {
4408 let file = write_parquet_from_iter(vec![(
4409 "non_utf8_binary",
4410 Arc::new(BinaryArray::from(vec![
4411 b"\xDE\x00\xFF".as_ref(),
4412 b"\xDE\x01\xAA".as_ref(),
4413 b"\xDE\x02\xFF".as_ref(),
4414 ])) as ArrayRef,
4415 )]);
4416 let supplied_fields = Fields::from(vec![Field::new(
4417 "non_utf8_binary",
4418 ArrowDataType::Utf8,
4419 false,
4420 )]);
4421
4422 let options = ArrowReaderOptions::new().with_schema(Arc::new(Schema::new(supplied_fields)));
4423 let mut arrow_reader = ParquetRecordBatchReaderBuilder::try_new_with_options(
4424 file.try_clone().unwrap(),
4425 options,
4426 )
4427 .expect("reader builder with schema")
4428 .build()
4429 .expect("reader with schema");
4430 arrow_reader.next().unwrap().unwrap_err();
4431 }
4432
4433 #[test]
4434 fn test_with_schema() {
4435 let nested_fields = Fields::from(vec![
4436 Field::new("utf8_to_dict", ArrowDataType::Utf8, false),
4437 Field::new("int64_to_ts_nano", ArrowDataType::Int64, false),
4438 ]);
4439
4440 let nested_arrays: Vec<ArrayRef> = vec![
4441 Arc::new(StringArray::from(vec!["a", "a", "a", "b"])) as ArrayRef,
4442 Arc::new(Int64Array::from(vec![1, 2, 3, 4])) as ArrayRef,
4443 ];
4444
4445 let nested = StructArray::try_new(nested_fields, nested_arrays, None).unwrap();
4446
4447 let file = write_parquet_from_iter(vec![
4448 (
4449 "int32_to_ts_second",
4450 Arc::new(Int32Array::from(vec![0, 1, 2, 3])) as ArrayRef,
4451 ),
4452 (
4453 "date32_to_date64",
4454 Arc::new(Date32Array::from(vec![0, 1, 2, 3])) as ArrayRef,
4455 ),
4456 ("nested", Arc::new(nested) as ArrayRef),
4457 ]);
4458
4459 let supplied_nested_fields = Fields::from(vec![
4460 Field::new(
4461 "utf8_to_dict",
4462 ArrowDataType::Dictionary(
4463 Box::new(ArrowDataType::Int32),
4464 Box::new(ArrowDataType::Utf8),
4465 ),
4466 false,
4467 ),
4468 Field::new(
4469 "int64_to_ts_nano",
4470 ArrowDataType::Timestamp(
4471 arrow::datatypes::TimeUnit::Nanosecond,
4472 Some("+10:00".into()),
4473 ),
4474 false,
4475 ),
4476 ]);
4477
4478 let supplied_schema = Arc::new(Schema::new(vec![
4479 Field::new(
4480 "int32_to_ts_second",
4481 ArrowDataType::Timestamp(arrow::datatypes::TimeUnit::Second, Some("+01:00".into())),
4482 false,
4483 ),
4484 Field::new("date32_to_date64", ArrowDataType::Date64, false),
4485 Field::new(
4486 "nested",
4487 ArrowDataType::Struct(supplied_nested_fields),
4488 false,
4489 ),
4490 ]));
4491
4492 let options = ArrowReaderOptions::new().with_schema(supplied_schema.clone());
4493 let mut arrow_reader = ParquetRecordBatchReaderBuilder::try_new_with_options(
4494 file.try_clone().unwrap(),
4495 options,
4496 )
4497 .expect("reader builder with schema")
4498 .build()
4499 .expect("reader with schema");
4500
4501 assert_eq!(arrow_reader.schema(), supplied_schema);
4502 let batch = arrow_reader.next().unwrap().unwrap();
4503 assert_eq!(batch.num_columns(), 3);
4504 assert_eq!(batch.num_rows(), 4);
4505 assert_eq!(
4506 batch
4507 .column(0)
4508 .as_any()
4509 .downcast_ref::<TimestampSecondArray>()
4510 .expect("downcast to timestamp second")
4511 .value_as_datetime_with_tz(0, "+01:00".parse().unwrap())
4512 .map(|v| v.to_string())
4513 .expect("value as datetime"),
4514 "1970-01-01 01:00:00 +01:00"
4515 );
4516 assert_eq!(
4517 batch
4518 .column(1)
4519 .as_any()
4520 .downcast_ref::<Date64Array>()
4521 .expect("downcast to date64")
4522 .value_as_date(0)
4523 .map(|v| v.to_string())
4524 .expect("value as date"),
4525 "1970-01-01"
4526 );
4527
4528 let nested = batch
4529 .column(2)
4530 .as_any()
4531 .downcast_ref::<StructArray>()
4532 .expect("downcast to struct");
4533
4534 let nested_dict = nested
4535 .column(0)
4536 .as_any()
4537 .downcast_ref::<Int32DictionaryArray>()
4538 .expect("downcast to dictionary");
4539
4540 assert_eq!(
4541 nested_dict
4542 .values()
4543 .as_any()
4544 .downcast_ref::<StringArray>()
4545 .expect("downcast to string")
4546 .iter()
4547 .collect::<Vec<_>>(),
4548 vec![Some("a"), Some("b")]
4549 );
4550
4551 assert_eq!(
4552 nested_dict.keys().iter().collect::<Vec<_>>(),
4553 vec![Some(0), Some(0), Some(0), Some(1)]
4554 );
4555
4556 assert_eq!(
4557 nested
4558 .column(1)
4559 .as_any()
4560 .downcast_ref::<TimestampNanosecondArray>()
4561 .expect("downcast to timestamp nanosecond")
4562 .value_as_datetime_with_tz(0, "+10:00".parse().unwrap())
4563 .map(|v| v.to_string())
4564 .expect("value as datetime"),
4565 "1970-01-01 10:00:00.000000001 +10:00"
4566 );
4567 }
4568
4569 #[test]
4570 fn test_empty_projection() {
4571 let testdata = arrow::util::test_util::parquet_test_data();
4572 let path = format!("{testdata}/alltypes_plain.parquet");
4573 let file = File::open(path).unwrap();
4574
4575 let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
4576 let file_metadata = builder.metadata().file_metadata();
4577 let expected_rows = file_metadata.num_rows() as usize;
4578
4579 let mask = ProjectionMask::leaves(builder.parquet_schema(), []);
4580 let batch_reader = builder
4581 .with_projection(mask)
4582 .with_batch_size(2)
4583 .build()
4584 .unwrap();
4585
4586 let mut total_rows = 0;
4587 for maybe_batch in batch_reader {
4588 let batch = maybe_batch.unwrap();
4589 total_rows += batch.num_rows();
4590 assert_eq!(batch.num_columns(), 0);
4591 assert!(batch.num_rows() <= 2);
4592 }
4593
4594 assert_eq!(total_rows, expected_rows);
4595 }
4596
4597 fn test_row_group_batch(row_group_size: usize, batch_size: usize) {
4598 let schema = Arc::new(Schema::new(vec![Field::new(
4599 "list",
4600 ArrowDataType::List(Arc::new(Field::new_list_field(ArrowDataType::Int32, true))),
4601 true,
4602 )]));
4603
4604 let mut buf = Vec::with_capacity(1024);
4605
4606 let mut writer = ArrowWriter::try_new(
4607 &mut buf,
4608 schema.clone(),
4609 Some(
4610 WriterProperties::builder()
4611 .set_max_row_group_row_count(Some(row_group_size))
4612 .build(),
4613 ),
4614 )
4615 .unwrap();
4616 for _ in 0..2 {
4617 let mut list_builder = ListBuilder::new(Int32Builder::with_capacity(batch_size));
4618 for _ in 0..(batch_size) {
4619 list_builder.append(true);
4620 }
4621 let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(list_builder.finish())])
4622 .unwrap();
4623 writer.write(&batch).unwrap();
4624 }
4625 writer.close().unwrap();
4626
4627 let mut record_reader =
4628 ParquetRecordBatchReader::try_new(Bytes::from(buf), batch_size).unwrap();
4629 assert_eq!(
4630 batch_size,
4631 record_reader.next().unwrap().unwrap().num_rows()
4632 );
4633 assert_eq!(
4634 batch_size,
4635 record_reader.next().unwrap().unwrap().num_rows()
4636 );
4637 }
4638
4639 #[test]
4640 fn test_row_group_exact_multiple() {
4641 const BATCH_SIZE: usize = REPETITION_LEVELS_BATCH_SIZE;
4642 test_row_group_batch(8, 8);
4643 test_row_group_batch(10, 8);
4644 test_row_group_batch(8, 10);
4645 test_row_group_batch(BATCH_SIZE, BATCH_SIZE);
4646 test_row_group_batch(BATCH_SIZE + 1, BATCH_SIZE);
4647 test_row_group_batch(BATCH_SIZE, BATCH_SIZE + 1);
4648 test_row_group_batch(BATCH_SIZE, BATCH_SIZE - 1);
4649 test_row_group_batch(BATCH_SIZE - 1, BATCH_SIZE);
4650 }
4651
4652 fn get_expected_batches(
4655 column: &RecordBatch,
4656 selection: &RowSelection,
4657 batch_size: usize,
4658 ) -> Vec<RecordBatch> {
4659 let mut expected_batches = vec![];
4660
4661 let mut selection: VecDeque<_> = selection.clone().into();
4662 let mut row_offset = 0;
4663 let mut last_start = None;
4664 while row_offset < column.num_rows() && !selection.is_empty() {
4665 let mut batch_remaining = batch_size.min(column.num_rows() - row_offset);
4666 while batch_remaining > 0 && !selection.is_empty() {
4667 let (to_read, skip) = match selection.front_mut() {
4668 Some(selection) if selection.row_count > batch_remaining => {
4669 selection.row_count -= batch_remaining;
4670 (batch_remaining, selection.skip)
4671 }
4672 Some(_) => {
4673 let select = selection.pop_front().unwrap();
4674 (select.row_count, select.skip)
4675 }
4676 None => break,
4677 };
4678
4679 batch_remaining -= to_read;
4680
4681 match skip {
4682 true => {
4683 if let Some(last_start) = last_start.take() {
4684 expected_batches.push(column.slice(last_start, row_offset - last_start))
4685 }
4686 row_offset += to_read
4687 }
4688 false => {
4689 last_start.get_or_insert(row_offset);
4690 row_offset += to_read
4691 }
4692 }
4693 }
4694 }
4695
4696 if let Some(last_start) = last_start.take() {
4697 expected_batches.push(column.slice(last_start, row_offset - last_start))
4698 }
4699
4700 for batch in &expected_batches[..expected_batches.len() - 1] {
4702 assert_eq!(batch.num_rows(), batch_size);
4703 }
4704
4705 expected_batches
4706 }
4707
4708 fn create_test_selection(
4709 step_len: usize,
4710 total_len: usize,
4711 skip_first: bool,
4712 ) -> (RowSelection, usize) {
4713 let mut remaining = total_len;
4714 let mut skip = skip_first;
4715 let mut vec = vec![];
4716 let mut selected_count = 0;
4717 while remaining != 0 {
4718 let step = if remaining > step_len {
4719 step_len
4720 } else {
4721 remaining
4722 };
4723 vec.push(RowSelector {
4724 row_count: step,
4725 skip,
4726 });
4727 remaining -= step;
4728 if !skip {
4729 selected_count += step;
4730 }
4731 skip = !skip;
4732 }
4733 (vec.into(), selected_count)
4734 }
4735
4736 #[test]
4737 fn test_scan_row_with_selection() {
4738 let testdata = arrow::util::test_util::parquet_test_data();
4739 let path = format!("{testdata}/alltypes_tiny_pages_plain.parquet");
4740 let test_file = File::open(&path).unwrap();
4741
4742 let mut serial_reader =
4743 ParquetRecordBatchReader::try_new(File::open(&path).unwrap(), 7300).unwrap();
4744 let data = serial_reader.next().unwrap().unwrap();
4745
4746 let do_test = |batch_size: usize, selection_len: usize| {
4747 for skip_first in [false, true] {
4748 let selections = create_test_selection(batch_size, data.num_rows(), skip_first).0;
4749
4750 let expected = get_expected_batches(&data, &selections, batch_size);
4751 let skip_reader = create_skip_reader(&test_file, batch_size, selections);
4752 assert_eq!(
4753 skip_reader.collect::<Result<Vec<_>, _>>().unwrap(),
4754 expected,
4755 "batch_size: {batch_size}, selection_len: {selection_len}, skip_first: {skip_first}"
4756 );
4757 }
4758 };
4759
4760 do_test(1000, 1000);
4763
4764 do_test(20, 20);
4766
4767 do_test(20, 5);
4769
4770 do_test(20, 5);
4773
4774 fn create_skip_reader(
4775 test_file: &File,
4776 batch_size: usize,
4777 selections: RowSelection,
4778 ) -> ParquetRecordBatchReader {
4779 let options =
4780 ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required);
4781 let file = test_file.try_clone().unwrap();
4782 ParquetRecordBatchReaderBuilder::try_new_with_options(file, options)
4783 .unwrap()
4784 .with_batch_size(batch_size)
4785 .with_row_selection(selections)
4786 .build()
4787 .unwrap()
4788 }
4789 }
4790
4791 #[test]
4792 fn test_batch_size_overallocate() {
4793 let testdata = arrow::util::test_util::parquet_test_data();
4794 let path = format!("{testdata}/alltypes_plain.parquet");
4796 let test_file = File::open(path).unwrap();
4797
4798 let builder = ParquetRecordBatchReaderBuilder::try_new(test_file).unwrap();
4799 let num_rows = builder.metadata.file_metadata().num_rows();
4800 let reader = builder
4801 .with_batch_size(1024)
4802 .with_projection(ProjectionMask::all())
4803 .build()
4804 .unwrap();
4805 assert_ne!(1024, num_rows);
4806 assert_eq!(reader.read_plan.batch_size(), num_rows as usize);
4807 }
4808
4809 #[test]
4810 fn test_read_with_page_index_enabled() {
4811 let testdata = arrow::util::test_util::parquet_test_data();
4812
4813 {
4814 let path = format!("{testdata}/alltypes_tiny_pages.parquet");
4816 let test_file = File::open(path).unwrap();
4817 let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(
4818 test_file,
4819 ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required),
4820 )
4821 .unwrap();
4822 assert!(!builder.metadata().offset_index().unwrap()[0].is_empty());
4823 let reader = builder.build().unwrap();
4824 let batches = reader.collect::<Result<Vec<_>, _>>().unwrap();
4825 assert_eq!(batches.len(), 8);
4826 }
4827
4828 {
4829 let path = format!("{testdata}/alltypes_plain.parquet");
4831 let test_file = File::open(path).unwrap();
4832 let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(
4833 test_file,
4834 ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Required),
4835 )
4836 .unwrap();
4837 assert!(builder.metadata().offset_index().is_none());
4840 let reader = builder.build().unwrap();
4841 let batches = reader.collect::<Result<Vec<_>, _>>().unwrap();
4842 assert_eq!(batches.len(), 1);
4843 }
4844 }
4845
4846 #[test]
4847 fn test_raw_repetition() {
4848 const MESSAGE_TYPE: &str = "
4849 message Log {
4850 OPTIONAL INT32 eventType;
4851 REPEATED INT32 category;
4852 REPEATED group filter {
4853 OPTIONAL INT32 error;
4854 }
4855 }
4856 ";
4857 let schema = Arc::new(parse_message_type(MESSAGE_TYPE).unwrap());
4858 let props = Default::default();
4859
4860 let mut buf = Vec::with_capacity(1024);
4861 let mut writer = SerializedFileWriter::new(&mut buf, schema, props).unwrap();
4862 let mut row_group_writer = writer.next_row_group().unwrap();
4863
4864 let mut col_writer = row_group_writer.next_column().unwrap().unwrap();
4866 col_writer
4867 .typed::<Int32Type>()
4868 .write_batch(&[1], Some(&[1]), None)
4869 .unwrap();
4870 col_writer.close().unwrap();
4871 let mut col_writer = row_group_writer.next_column().unwrap().unwrap();
4873 col_writer
4874 .typed::<Int32Type>()
4875 .write_batch(&[1, 1], Some(&[1, 1]), Some(&[0, 1]))
4876 .unwrap();
4877 col_writer.close().unwrap();
4878 let mut col_writer = row_group_writer.next_column().unwrap().unwrap();
4880 col_writer
4881 .typed::<Int32Type>()
4882 .write_batch(&[1], Some(&[1]), Some(&[0]))
4883 .unwrap();
4884 col_writer.close().unwrap();
4885
4886 let rg_md = row_group_writer.close().unwrap();
4887 assert_eq!(rg_md.num_rows(), 1);
4888 writer.close().unwrap();
4889
4890 let bytes = Bytes::from(buf);
4891
4892 let mut no_mask = ParquetRecordBatchReader::try_new(bytes.clone(), 1024).unwrap();
4893 let full = no_mask.next().unwrap().unwrap();
4894
4895 assert_eq!(full.num_columns(), 3);
4896
4897 for idx in 0..3 {
4898 let b = ParquetRecordBatchReaderBuilder::try_new(bytes.clone()).unwrap();
4899 let mask = ProjectionMask::leaves(b.parquet_schema(), [idx]);
4900 let mut reader = b.with_projection(mask).build().unwrap();
4901 let projected = reader.next().unwrap().unwrap();
4902
4903 assert_eq!(projected.num_columns(), 1);
4904 assert_eq!(full.column(idx), projected.column(0));
4905 }
4906 }
4907
4908 #[test]
4909 fn test_read_lz4_raw() {
4910 let testdata = arrow::util::test_util::parquet_test_data();
4911 let path = format!("{testdata}/lz4_raw_compressed.parquet");
4912 let file = File::open(path).unwrap();
4913
4914 let batches = ParquetRecordBatchReader::try_new(file, 1024)
4915 .unwrap()
4916 .collect::<Result<Vec<_>, _>>()
4917 .unwrap();
4918 assert_eq!(batches.len(), 1);
4919 let batch = &batches[0];
4920
4921 assert_eq!(batch.num_columns(), 3);
4922 assert_eq!(batch.num_rows(), 4);
4923
4924 let a: &Int64Array = batch.column(0).as_any().downcast_ref().unwrap();
4926 assert_eq!(
4927 a.values(),
4928 &[1593604800, 1593604800, 1593604801, 1593604801]
4929 );
4930
4931 let a: &BinaryArray = batch.column(1).as_any().downcast_ref().unwrap();
4932 let a: Vec<_> = a.iter().flatten().collect();
4933 assert_eq!(a, &[b"abc", b"def", b"abc", b"def"]);
4934
4935 let a: &Float64Array = batch.column(2).as_any().downcast_ref().unwrap();
4936 assert_eq!(a.values(), &[42.000000, 7.700000, 42.125000, 7.700000]);
4937 }
4938
4939 #[test]
4949 fn test_read_lz4_hadoop_fallback() {
4950 for file in [
4951 "hadoop_lz4_compressed.parquet",
4952 "non_hadoop_lz4_compressed.parquet",
4953 ] {
4954 let testdata = arrow::util::test_util::parquet_test_data();
4955 let path = format!("{testdata}/{file}");
4956 let file = File::open(path).unwrap();
4957 let expected_rows = 4;
4958
4959 let batches = ParquetRecordBatchReader::try_new(file, expected_rows)
4960 .unwrap()
4961 .collect::<Result<Vec<_>, _>>()
4962 .unwrap();
4963 assert_eq!(batches.len(), 1);
4964 let batch = &batches[0];
4965
4966 assert_eq!(batch.num_columns(), 3);
4967 assert_eq!(batch.num_rows(), expected_rows);
4968
4969 let a: &Int64Array = batch.column(0).as_any().downcast_ref().unwrap();
4970 assert_eq!(
4971 a.values(),
4972 &[1593604800, 1593604800, 1593604801, 1593604801]
4973 );
4974
4975 let b: &BinaryArray = batch.column(1).as_any().downcast_ref().unwrap();
4976 let b: Vec<_> = b.iter().flatten().collect();
4977 assert_eq!(b, &[b"abc", b"def", b"abc", b"def"]);
4978
4979 let c: &Float64Array = batch.column(2).as_any().downcast_ref().unwrap();
4980 assert_eq!(c.values(), &[42.0, 7.7, 42.125, 7.7]);
4981 }
4982 }
4983
4984 #[test]
4985 fn test_read_lz4_hadoop_large() {
4986 let testdata = arrow::util::test_util::parquet_test_data();
4987 let path = format!("{testdata}/hadoop_lz4_compressed_larger.parquet");
4988 let file = File::open(path).unwrap();
4989 let expected_rows = 10000;
4990
4991 let batches = ParquetRecordBatchReader::try_new(file, expected_rows)
4992 .unwrap()
4993 .collect::<Result<Vec<_>, _>>()
4994 .unwrap();
4995 assert_eq!(batches.len(), 1);
4996 let batch = &batches[0];
4997
4998 assert_eq!(batch.num_columns(), 1);
4999 assert_eq!(batch.num_rows(), expected_rows);
5000
5001 let a: &StringArray = batch.column(0).as_any().downcast_ref().unwrap();
5002 let a: Vec<_> = a.iter().flatten().collect();
5003 assert_eq!(a[0], "c7ce6bef-d5b0-4863-b199-8ea8c7fb117b");
5004 assert_eq!(a[1], "e8fb9197-cb9f-4118-b67f-fbfa65f61843");
5005 assert_eq!(a[expected_rows - 2], "ab52a0cc-c6bb-4d61-8a8f-166dc4b8b13c");
5006 assert_eq!(a[expected_rows - 1], "85440778-460a-41ac-aa2e-ac3ee41696bf");
5007 }
5008
5009 #[test]
5010 #[cfg(feature = "snap")]
5011 fn test_read_nested_lists() {
5012 let testdata = arrow::util::test_util::parquet_test_data();
5013 let path = format!("{testdata}/nested_lists.snappy.parquet");
5014 let file = File::open(path).unwrap();
5015
5016 let f = file.try_clone().unwrap();
5017 let mut reader = ParquetRecordBatchReader::try_new(f, 60).unwrap();
5018 let expected = reader.next().unwrap().unwrap();
5019 assert_eq!(expected.num_rows(), 3);
5020
5021 let selection = RowSelection::from(vec![
5022 RowSelector::skip(1),
5023 RowSelector::select(1),
5024 RowSelector::skip(1),
5025 ]);
5026 let mut reader = ParquetRecordBatchReaderBuilder::try_new(file)
5027 .unwrap()
5028 .with_row_selection(selection)
5029 .build()
5030 .unwrap();
5031
5032 let actual = reader.next().unwrap().unwrap();
5033 assert_eq!(actual.num_rows(), 1);
5034 assert_eq!(actual.column(0), &expected.column(0).slice(1, 1));
5035 }
5036
5037 #[test]
5038 fn test_arbitrary_decimal() {
5039 let values = [1, 2, 3, 4, 5, 6, 7, 8];
5040 let decimals_19_0 = Decimal128Array::from_iter_values(values)
5041 .with_precision_and_scale(19, 0)
5042 .unwrap();
5043 let decimals_12_0 = Decimal128Array::from_iter_values(values)
5044 .with_precision_and_scale(12, 0)
5045 .unwrap();
5046 let decimals_17_10 = Decimal128Array::from_iter_values(values)
5047 .with_precision_and_scale(17, 10)
5048 .unwrap();
5049
5050 let written = RecordBatch::try_from_iter([
5051 ("decimal_values_19_0", Arc::new(decimals_19_0) as ArrayRef),
5052 ("decimal_values_12_0", Arc::new(decimals_12_0) as ArrayRef),
5053 ("decimal_values_17_10", Arc::new(decimals_17_10) as ArrayRef),
5054 ])
5055 .unwrap();
5056
5057 let mut buffer = Vec::with_capacity(1024);
5058 let mut writer = ArrowWriter::try_new(&mut buffer, written.schema(), None).unwrap();
5059 writer.write(&written).unwrap();
5060 writer.close().unwrap();
5061
5062 let read = ParquetRecordBatchReader::try_new(Bytes::from(buffer), 8)
5063 .unwrap()
5064 .collect::<Result<Vec<_>, _>>()
5065 .unwrap();
5066
5067 assert_eq!(&written.slice(0, 8), &read[0]);
5068 }
5069
5070 #[test]
5071 fn test_list_skip() {
5072 let mut list = ListBuilder::new(Int32Builder::new());
5073 list.append_value([Some(1), Some(2)]);
5074 list.append_value([Some(3)]);
5075 list.append_value([Some(4)]);
5076 let list = list.finish();
5077 let batch = RecordBatch::try_from_iter([("l", Arc::new(list) as _)]).unwrap();
5078
5079 let props = WriterProperties::builder()
5081 .set_data_page_row_count_limit(1)
5082 .set_write_batch_size(2)
5083 .build();
5084
5085 let mut buffer = Vec::with_capacity(1024);
5086 let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema(), Some(props)).unwrap();
5087 writer.write(&batch).unwrap();
5088 writer.close().unwrap();
5089
5090 let selection = vec![RowSelector::skip(2), RowSelector::select(1)];
5091 let mut reader = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(buffer))
5092 .unwrap()
5093 .with_row_selection(selection.into())
5094 .build()
5095 .unwrap();
5096 let out = reader.next().unwrap().unwrap();
5097 assert_eq!(out.num_rows(), 1);
5098 assert_eq!(out, batch.slice(2, 1));
5099 }
5100
5101 fn test_decimal32_roundtrip() {
5102 let d = |values: Vec<i32>, p: u8| {
5103 let iter = values.into_iter();
5104 PrimitiveArray::<Decimal32Type>::from_iter_values(iter)
5105 .with_precision_and_scale(p, 2)
5106 .unwrap()
5107 };
5108
5109 let d1 = d(vec![1, 2, 3, 4, 5], 9);
5110 let batch = RecordBatch::try_from_iter([("d1", Arc::new(d1) as ArrayRef)]).unwrap();
5111
5112 let mut buffer = Vec::with_capacity(1024);
5113 let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema(), None).unwrap();
5114 writer.write(&batch).unwrap();
5115 writer.close().unwrap();
5116
5117 let builder = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(buffer)).unwrap();
5118 let t1 = builder.parquet_schema().columns()[0].physical_type();
5119 assert_eq!(t1, PhysicalType::INT32);
5120
5121 let mut reader = builder.build().unwrap();
5122 assert_eq!(batch.schema(), reader.schema());
5123
5124 let out = reader.next().unwrap().unwrap();
5125 assert_eq!(batch, out);
5126 }
5127
5128 fn test_decimal64_roundtrip() {
5129 let d = |values: Vec<i64>, p: u8| {
5133 let iter = values.into_iter();
5134 PrimitiveArray::<Decimal64Type>::from_iter_values(iter)
5135 .with_precision_and_scale(p, 2)
5136 .unwrap()
5137 };
5138
5139 let d1 = d(vec![1, 2, 3, 4, 5], 9);
5140 let d2 = d(vec![1, 2, 3, 4, 10.pow(10) - 1], 10);
5141 let d3 = d(vec![1, 2, 3, 4, 10.pow(18) - 1], 18);
5142
5143 let batch = RecordBatch::try_from_iter([
5144 ("d1", Arc::new(d1) as ArrayRef),
5145 ("d2", Arc::new(d2) as ArrayRef),
5146 ("d3", Arc::new(d3) as ArrayRef),
5147 ])
5148 .unwrap();
5149
5150 let mut buffer = Vec::with_capacity(1024);
5151 let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema(), None).unwrap();
5152 writer.write(&batch).unwrap();
5153 writer.close().unwrap();
5154
5155 let builder = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(buffer)).unwrap();
5156 let t1 = builder.parquet_schema().columns()[0].physical_type();
5157 assert_eq!(t1, PhysicalType::INT32);
5158 let t2 = builder.parquet_schema().columns()[1].physical_type();
5159 assert_eq!(t2, PhysicalType::INT64);
5160 let t3 = builder.parquet_schema().columns()[2].physical_type();
5161 assert_eq!(t3, PhysicalType::INT64);
5162
5163 let mut reader = builder.build().unwrap();
5164 assert_eq!(batch.schema(), reader.schema());
5165
5166 let out = reader.next().unwrap().unwrap();
5167 assert_eq!(batch, out);
5168 }
5169
5170 fn test_decimal_roundtrip<T: DecimalType>() {
5171 let d = |values: Vec<usize>, p: u8| {
5176 let iter = values.into_iter().map(T::Native::usize_as);
5177 PrimitiveArray::<T>::from_iter_values(iter)
5178 .with_precision_and_scale(p, 2)
5179 .unwrap()
5180 };
5181
5182 let d1 = d(vec![1, 2, 3, 4, 5], 9);
5183 let d2 = d(vec![1, 2, 3, 4, 10.pow(10) - 1], 10);
5184 let d3 = d(vec![1, 2, 3, 4, 10.pow(18) - 1], 18);
5185 let d4 = d(vec![1, 2, 3, 4, 10.pow(19) - 1], 19);
5186
5187 let batch = RecordBatch::try_from_iter([
5188 ("d1", Arc::new(d1) as ArrayRef),
5189 ("d2", Arc::new(d2) as ArrayRef),
5190 ("d3", Arc::new(d3) as ArrayRef),
5191 ("d4", Arc::new(d4) as ArrayRef),
5192 ])
5193 .unwrap();
5194
5195 let mut buffer = Vec::with_capacity(1024);
5196 let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema(), None).unwrap();
5197 writer.write(&batch).unwrap();
5198 writer.close().unwrap();
5199
5200 let builder = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(buffer)).unwrap();
5201 let t1 = builder.parquet_schema().columns()[0].physical_type();
5202 assert_eq!(t1, PhysicalType::INT32);
5203 let t2 = builder.parquet_schema().columns()[1].physical_type();
5204 assert_eq!(t2, PhysicalType::INT64);
5205 let t3 = builder.parquet_schema().columns()[2].physical_type();
5206 assert_eq!(t3, PhysicalType::INT64);
5207 let t4 = builder.parquet_schema().columns()[3].physical_type();
5208 assert_eq!(t4, PhysicalType::FIXED_LEN_BYTE_ARRAY);
5209
5210 let mut reader = builder.build().unwrap();
5211 assert_eq!(batch.schema(), reader.schema());
5212
5213 let out = reader.next().unwrap().unwrap();
5214 assert_eq!(batch, out);
5215 }
5216
5217 #[test]
5218 fn test_decimal() {
5219 test_decimal32_roundtrip();
5220 test_decimal64_roundtrip();
5221 test_decimal_roundtrip::<Decimal128Type>();
5222 test_decimal_roundtrip::<Decimal256Type>();
5223 }
5224
5225 #[test]
5226 fn test_list_selection() {
5227 let schema = Arc::new(Schema::new(vec![Field::new_list(
5228 "list",
5229 Field::new_list_field(ArrowDataType::Utf8, true),
5230 false,
5231 )]));
5232 let mut buf = Vec::with_capacity(1024);
5233
5234 let mut writer = ArrowWriter::try_new(&mut buf, schema.clone(), None).unwrap();
5235
5236 for i in 0..2 {
5237 let mut list_a_builder = ListBuilder::new(StringBuilder::new());
5238 for j in 0..1024 {
5239 list_a_builder.values().append_value(format!("{i} {j}"));
5240 list_a_builder.append(true);
5241 }
5242 let batch =
5243 RecordBatch::try_new(schema.clone(), vec![Arc::new(list_a_builder.finish())])
5244 .unwrap();
5245 writer.write(&batch).unwrap();
5246 }
5247 let _metadata = writer.close().unwrap();
5248
5249 let buf = Bytes::from(buf);
5250 let reader = ParquetRecordBatchReaderBuilder::try_new(buf)
5251 .unwrap()
5252 .with_row_selection(RowSelection::from(vec![
5253 RowSelector::skip(100),
5254 RowSelector::select(924),
5255 RowSelector::skip(100),
5256 RowSelector::select(924),
5257 ]))
5258 .build()
5259 .unwrap();
5260
5261 let batches = reader.collect::<Result<Vec<_>, _>>().unwrap();
5262 let batch = concat_batches(&schema, &batches).unwrap();
5263
5264 assert_eq!(batch.num_rows(), 924 * 2);
5265 let list = batch.column(0).as_list::<i32>();
5266
5267 for w in list.value_offsets().windows(2) {
5268 assert_eq!(w[0] + 1, w[1])
5269 }
5270 let mut values = list.values().as_string::<i32>().iter();
5271
5272 for i in 0..2 {
5273 for j in 100..1024 {
5274 let expected = format!("{i} {j}");
5275 assert_eq!(values.next().unwrap().unwrap(), &expected);
5276 }
5277 }
5278 }
5279
5280 #[test]
5281 fn test_list_selection_fuzz() {
5282 let mut rng = rng();
5283 let schema = Arc::new(Schema::new(vec![Field::new_list(
5284 "list",
5285 Field::new_list(
5286 Field::LIST_FIELD_DEFAULT_NAME,
5287 Field::new_list_field(ArrowDataType::Int32, true),
5288 true,
5289 ),
5290 true,
5291 )]));
5292 let mut buf = Vec::with_capacity(1024);
5293 let mut writer = ArrowWriter::try_new(&mut buf, schema.clone(), None).unwrap();
5294
5295 let mut list_a_builder = ListBuilder::new(ListBuilder::new(Int32Builder::new()));
5296
5297 for _ in 0..2048 {
5298 if rng.random_bool(0.2) {
5299 list_a_builder.append(false);
5300 continue;
5301 }
5302
5303 let list_a_len = rng.random_range(0..10);
5304 let list_b_builder = list_a_builder.values();
5305
5306 for _ in 0..list_a_len {
5307 if rng.random_bool(0.2) {
5308 list_b_builder.append(false);
5309 continue;
5310 }
5311
5312 let list_b_len = rng.random_range(0..10);
5313 let int_builder = list_b_builder.values();
5314 for _ in 0..list_b_len {
5315 match rng.random_bool(0.2) {
5316 true => int_builder.append_null(),
5317 false => int_builder.append_value(rng.random()),
5318 }
5319 }
5320 list_b_builder.append(true)
5321 }
5322 list_a_builder.append(true);
5323 }
5324
5325 let array = Arc::new(list_a_builder.finish());
5326 let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
5327
5328 writer.write(&batch).unwrap();
5329 let _metadata = writer.close().unwrap();
5330
5331 let buf = Bytes::from(buf);
5332
5333 let cases = [
5334 vec![
5335 RowSelector::skip(100),
5336 RowSelector::select(924),
5337 RowSelector::skip(100),
5338 RowSelector::select(924),
5339 ],
5340 vec![
5341 RowSelector::select(924),
5342 RowSelector::skip(100),
5343 RowSelector::select(924),
5344 RowSelector::skip(100),
5345 ],
5346 vec![
5347 RowSelector::skip(1023),
5348 RowSelector::select(1),
5349 RowSelector::skip(1023),
5350 RowSelector::select(1),
5351 ],
5352 vec![
5353 RowSelector::select(1),
5354 RowSelector::skip(1023),
5355 RowSelector::select(1),
5356 RowSelector::skip(1023),
5357 ],
5358 ];
5359
5360 for batch_size in [100, 1024, 2048] {
5361 for selection in &cases {
5362 let selection = RowSelection::from(selection.clone());
5363 let reader = ParquetRecordBatchReaderBuilder::try_new(buf.clone())
5364 .unwrap()
5365 .with_row_selection(selection.clone())
5366 .with_batch_size(batch_size)
5367 .build()
5368 .unwrap();
5369
5370 let batches = reader.collect::<Result<Vec<_>, _>>().unwrap();
5371 let actual = concat_batches(batch.schema_ref(), &batches).unwrap();
5372 assert_eq!(actual.num_rows(), selection.row_count());
5373
5374 let mut batch_offset = 0;
5375 let mut actual_offset = 0;
5376 for selector in selection.iter() {
5377 if selector.skip {
5378 batch_offset += selector.row_count;
5379 continue;
5380 }
5381
5382 assert_eq!(
5383 batch.slice(batch_offset, selector.row_count),
5384 actual.slice(actual_offset, selector.row_count)
5385 );
5386
5387 batch_offset += selector.row_count;
5388 actual_offset += selector.row_count;
5389 }
5390 }
5391 }
5392 }
5393
5394 #[test]
5395 fn test_read_old_nested_list() {
5396 use arrow::datatypes::DataType;
5397 use arrow::datatypes::ToByteSlice;
5398
5399 let testdata = arrow::util::test_util::parquet_test_data();
5400 let path = format!("{testdata}/old_list_structure.parquet");
5409 let test_file = File::open(path).unwrap();
5410
5411 let a_values = Int32Array::from(vec![1, 2, 3, 4]);
5413
5414 let a_value_offsets = arrow::buffer::Buffer::from([0, 2, 4].to_byte_slice());
5416
5417 let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new(
5419 "array",
5420 DataType::Int32,
5421 false,
5422 ))))
5423 .len(2)
5424 .add_buffer(a_value_offsets)
5425 .add_child_data(a_values.into_data())
5426 .build()
5427 .unwrap();
5428 let a = ListArray::from(a_list_data);
5429
5430 let builder = ParquetRecordBatchReaderBuilder::try_new(test_file).unwrap();
5431 let mut reader = builder.build().unwrap();
5432 let out = reader.next().unwrap().unwrap();
5433 assert_eq!(out.num_rows(), 1);
5434 assert_eq!(out.num_columns(), 1);
5435 let c0 = out.column(0);
5437 let c0arr = c0.as_any().downcast_ref::<ListArray>().unwrap();
5438 let r0 = c0arr.value(0);
5440 let r0arr = r0.as_any().downcast_ref::<ListArray>().unwrap();
5441 assert_eq!(r0arr, &a);
5442 }
5443
5444 #[test]
5445 fn test_map_no_value() {
5446 let testdata = arrow::util::test_util::parquet_test_data();
5466 let path = format!("{testdata}/map_no_value.parquet");
5467 let file = File::open(path).unwrap();
5468
5469 let mut reader = ParquetRecordBatchReaderBuilder::try_new(file)
5470 .unwrap()
5471 .build()
5472 .unwrap();
5473 let out = reader.next().unwrap().unwrap();
5474 assert_eq!(out.num_rows(), 3);
5475 assert_eq!(out.num_columns(), 3);
5476 let c0 = out.column(1).as_list::<i32>();
5478 let c1 = out.column(2).as_list::<i32>();
5479 assert_eq!(c0.len(), c1.len());
5480 c0.iter().zip(c1.iter()).for_each(|(l, r)| assert_eq!(l, r));
5481 }
5482
5483 #[test]
5484 fn test_read_unknown_logical_type() {
5485 let testdata = arrow::util::test_util::parquet_test_data();
5486 let path = format!("{testdata}/unknown-logical-type.parquet");
5487 let test_file = File::open(path).unwrap();
5488
5489 let builder = ParquetRecordBatchReaderBuilder::try_new(test_file)
5490 .expect("Error creating reader builder");
5491
5492 let schema = builder.metadata().file_metadata().schema_descr();
5493 assert_eq!(
5494 schema.column(0).logical_type_ref(),
5495 Some(&LogicalType::String)
5496 );
5497 assert_eq!(
5498 schema.column(1).logical_type_ref(),
5499 Some(&LogicalType::_Unknown { field_id: 2555 })
5500 );
5501 assert_eq!(schema.column(1).physical_type(), PhysicalType::BYTE_ARRAY);
5502
5503 let mut reader = builder.build().unwrap();
5504 let out = reader.next().unwrap().unwrap();
5505 assert_eq!(out.num_rows(), 3);
5506 assert_eq!(out.num_columns(), 2);
5507 }
5508
5509 #[test]
5510 fn test_read_row_numbers() {
5511 let file = write_parquet_from_iter(vec![(
5512 "value",
5513 Arc::new(Int64Array::from(vec![1, 2, 3])) as ArrayRef,
5514 )]);
5515 let supplied_fields = Fields::from(vec![Field::new("value", ArrowDataType::Int64, false)]);
5516
5517 let row_number_field = Arc::new(
5518 Field::new("row_number", ArrowDataType::Int64, false).with_extension_type(RowNumber),
5519 );
5520
5521 let options = ArrowReaderOptions::new()
5522 .with_schema(Arc::new(Schema::new(supplied_fields)))
5523 .with_virtual_columns(vec![row_number_field.clone()])
5524 .unwrap();
5525 let mut arrow_reader = ParquetRecordBatchReaderBuilder::try_new_with_options(
5526 file.try_clone().unwrap(),
5527 options,
5528 )
5529 .expect("reader builder with schema")
5530 .build()
5531 .expect("reader with schema");
5532
5533 let batch = arrow_reader.next().unwrap().unwrap();
5534 let schema = Arc::new(Schema::new(vec![
5535 Field::new("value", ArrowDataType::Int64, false),
5536 (*row_number_field).clone(),
5537 ]));
5538
5539 assert_eq!(batch.schema(), schema);
5540 assert_eq!(batch.num_columns(), 2);
5541 assert_eq!(batch.num_rows(), 3);
5542 assert_eq!(
5543 batch
5544 .column(0)
5545 .as_primitive::<types::Int64Type>()
5546 .iter()
5547 .collect::<Vec<_>>(),
5548 vec![Some(1), Some(2), Some(3)]
5549 );
5550 assert_eq!(
5551 batch
5552 .column(1)
5553 .as_primitive::<types::Int64Type>()
5554 .iter()
5555 .collect::<Vec<_>>(),
5556 vec![Some(0), Some(1), Some(2)]
5557 );
5558 }
5559
5560 #[test]
5561 fn test_read_only_row_numbers() {
5562 let file = write_parquet_from_iter(vec![(
5563 "value",
5564 Arc::new(Int64Array::from(vec![1, 2, 3])) as ArrayRef,
5565 )]);
5566 let row_number_field = Arc::new(
5567 Field::new("row_number", ArrowDataType::Int64, false).with_extension_type(RowNumber),
5568 );
5569 let options = ArrowReaderOptions::new()
5570 .with_virtual_columns(vec![row_number_field.clone()])
5571 .unwrap();
5572 let metadata = ArrowReaderMetadata::load(&file, options).unwrap();
5573 let num_columns = metadata
5574 .metadata
5575 .file_metadata()
5576 .schema_descr()
5577 .num_columns();
5578
5579 let mut arrow_reader = ParquetRecordBatchReaderBuilder::new_with_metadata(file, metadata)
5580 .with_projection(ProjectionMask::none(num_columns))
5581 .build()
5582 .expect("reader with schema");
5583
5584 let batch = arrow_reader.next().unwrap().unwrap();
5585 let schema = Arc::new(Schema::new(vec![row_number_field]));
5586
5587 assert_eq!(batch.schema(), schema);
5588 assert_eq!(batch.num_columns(), 1);
5589 assert_eq!(batch.num_rows(), 3);
5590 assert_eq!(
5591 batch
5592 .column(0)
5593 .as_primitive::<types::Int64Type>()
5594 .iter()
5595 .collect::<Vec<_>>(),
5596 vec![Some(0), Some(1), Some(2)]
5597 );
5598 }
5599
5600 #[test]
5601 fn test_read_row_numbers_row_group_order() -> Result<()> {
5602 let array = Int64Array::from_iter_values(5000..5100);
5604 let batch = RecordBatch::try_from_iter([("col", Arc::new(array) as ArrayRef)])?;
5605 let mut buffer = Vec::new();
5606 let options = WriterProperties::builder()
5607 .set_max_row_group_row_count(Some(50))
5608 .build();
5609 let mut writer = ArrowWriter::try_new(&mut buffer, batch.schema().clone(), Some(options))?;
5610 for batch_chunk in (0..10).map(|i| batch.slice(i * 10, 10)) {
5612 writer.write(&batch_chunk)?;
5613 }
5614 writer.close()?;
5615
5616 let row_number_field = Arc::new(
5617 Field::new("row_number", ArrowDataType::Int64, false).with_extension_type(RowNumber),
5618 );
5619
5620 let buffer = Bytes::from(buffer);
5621
5622 let options =
5623 ArrowReaderOptions::new().with_virtual_columns(vec![row_number_field.clone()])?;
5624
5625 let arrow_reader =
5627 ParquetRecordBatchReaderBuilder::try_new_with_options(buffer.clone(), options.clone())?
5628 .build()?;
5629
5630 assert_eq!(
5631 ValuesAndRowNumbers {
5632 values: (5000..5100).collect(),
5633 row_numbers: (0..100).collect()
5634 },
5635 ValuesAndRowNumbers::new_from_reader(arrow_reader)
5636 );
5637
5638 let arrow_reader = ParquetRecordBatchReaderBuilder::try_new_with_options(buffer, options)?
5640 .with_row_groups(vec![1, 0])
5641 .build()?;
5642
5643 assert_eq!(
5644 ValuesAndRowNumbers {
5645 values: (5050..5100).chain(5000..5050).collect(),
5646 row_numbers: (50..100).chain(0..50).collect(),
5647 },
5648 ValuesAndRowNumbers::new_from_reader(arrow_reader)
5649 );
5650
5651 Ok(())
5652 }
5653
5654 #[derive(Debug, PartialEq)]
5655 struct ValuesAndRowNumbers {
5656 values: Vec<i64>,
5657 row_numbers: Vec<i64>,
5658 }
5659 impl ValuesAndRowNumbers {
5660 fn new_from_reader(reader: ParquetRecordBatchReader) -> Self {
5661 let mut values = vec![];
5662 let mut row_numbers = vec![];
5663 for batch in reader {
5664 let batch = batch.expect("Could not read batch");
5665 values.extend(
5666 batch
5667 .column_by_name("col")
5668 .expect("Could not get col column")
5669 .as_primitive::<arrow::datatypes::Int64Type>()
5670 .iter()
5671 .map(|v| v.expect("Could not get value")),
5672 );
5673
5674 row_numbers.extend(
5675 batch
5676 .column_by_name("row_number")
5677 .expect("Could not get row_number column")
5678 .as_primitive::<arrow::datatypes::Int64Type>()
5679 .iter()
5680 .map(|v| v.expect("Could not get row number"))
5681 .collect::<Vec<_>>(),
5682 );
5683 }
5684 Self {
5685 values,
5686 row_numbers,
5687 }
5688 }
5689 }
5690
5691 #[test]
5692 fn test_with_virtual_columns_rejects_non_virtual_fields() {
5693 let regular_field = Arc::new(Field::new("regular_column", ArrowDataType::Int64, false));
5695 assert_eq!(
5696 ArrowReaderOptions::new()
5697 .with_virtual_columns(vec![regular_field])
5698 .unwrap_err()
5699 .to_string(),
5700 "Parquet error: Field 'regular_column' is not a virtual column. Virtual columns must have extension type names starting with 'arrow.virtual.'"
5701 );
5702 }
5703
5704 #[test]
5705 fn test_row_numbers_with_multiple_row_groups() {
5706 test_row_numbers_with_multiple_row_groups_helper(
5707 false,
5708 |path, selection, _row_filter, batch_size| {
5709 let file = File::open(path).unwrap();
5710 let row_number_field = Arc::new(
5711 Field::new("row_number", ArrowDataType::Int64, false)
5712 .with_extension_type(RowNumber),
5713 );
5714 let options = ArrowReaderOptions::new()
5715 .with_virtual_columns(vec![row_number_field])
5716 .unwrap();
5717 let reader = ParquetRecordBatchReaderBuilder::try_new_with_options(file, options)
5718 .unwrap()
5719 .with_row_selection(selection)
5720 .with_batch_size(batch_size)
5721 .build()
5722 .expect("Could not create reader");
5723 reader
5724 .collect::<Result<Vec<_>, _>>()
5725 .expect("Could not read")
5726 },
5727 );
5728 }
5729
5730 #[test]
5731 fn test_row_numbers_with_multiple_row_groups_and_filter() {
5732 test_row_numbers_with_multiple_row_groups_helper(
5733 true,
5734 |path, selection, row_filter, batch_size| {
5735 let file = File::open(path).unwrap();
5736 let row_number_field = Arc::new(
5737 Field::new("row_number", ArrowDataType::Int64, false)
5738 .with_extension_type(RowNumber),
5739 );
5740 let options = ArrowReaderOptions::new()
5741 .with_virtual_columns(vec![row_number_field])
5742 .unwrap();
5743 let reader = ParquetRecordBatchReaderBuilder::try_new_with_options(file, options)
5744 .unwrap()
5745 .with_row_selection(selection)
5746 .with_batch_size(batch_size)
5747 .with_row_filter(row_filter.expect("No filter"))
5748 .build()
5749 .expect("Could not create reader");
5750 reader
5751 .collect::<Result<Vec<_>, _>>()
5752 .expect("Could not read")
5753 },
5754 );
5755 }
5756
5757 #[test]
5758 fn test_read_row_group_indices() {
5759 let array1 = Int64Array::from(vec![1, 2]);
5761 let array2 = Int64Array::from(vec![3, 4]);
5762 let array3 = Int64Array::from(vec![5, 6]);
5763
5764 let batch1 =
5765 RecordBatch::try_from_iter(vec![("value", Arc::new(array1) as ArrayRef)]).unwrap();
5766 let batch2 =
5767 RecordBatch::try_from_iter(vec![("value", Arc::new(array2) as ArrayRef)]).unwrap();
5768 let batch3 =
5769 RecordBatch::try_from_iter(vec![("value", Arc::new(array3) as ArrayRef)]).unwrap();
5770
5771 let mut buffer = Vec::new();
5772 let options = WriterProperties::builder()
5773 .set_max_row_group_row_count(Some(2))
5774 .build();
5775 let mut writer = ArrowWriter::try_new(&mut buffer, batch1.schema(), Some(options)).unwrap();
5776 writer.write(&batch1).unwrap();
5777 writer.write(&batch2).unwrap();
5778 writer.write(&batch3).unwrap();
5779 writer.close().unwrap();
5780
5781 let file = Bytes::from(buffer);
5782 let row_group_index_field = Arc::new(
5783 Field::new("row_group_index", ArrowDataType::Int64, false)
5784 .with_extension_type(RowGroupIndex),
5785 );
5786
5787 let options = ArrowReaderOptions::new()
5788 .with_virtual_columns(vec![row_group_index_field.clone()])
5789 .unwrap();
5790 let mut arrow_reader =
5791 ParquetRecordBatchReaderBuilder::try_new_with_options(file.clone(), options)
5792 .expect("reader builder with virtual columns")
5793 .build()
5794 .expect("reader with virtual columns");
5795
5796 let batch = arrow_reader.next().unwrap().unwrap();
5797
5798 assert_eq!(batch.num_columns(), 2);
5799 assert_eq!(batch.num_rows(), 6);
5800
5801 assert_eq!(
5802 batch
5803 .column(0)
5804 .as_primitive::<types::Int64Type>()
5805 .iter()
5806 .collect::<Vec<_>>(),
5807 vec![Some(1), Some(2), Some(3), Some(4), Some(5), Some(6)]
5808 );
5809
5810 assert_eq!(
5811 batch
5812 .column(1)
5813 .as_primitive::<types::Int64Type>()
5814 .iter()
5815 .collect::<Vec<_>>(),
5816 vec![Some(0), Some(0), Some(1), Some(1), Some(2), Some(2)]
5817 );
5818 }
5819
5820 #[test]
5821 fn test_read_only_row_group_indices() {
5822 let array1 = Int64Array::from(vec![1, 2, 3]);
5823 let array2 = Int64Array::from(vec![4, 5]);
5824
5825 let batch1 =
5826 RecordBatch::try_from_iter(vec![("value", Arc::new(array1) as ArrayRef)]).unwrap();
5827 let batch2 =
5828 RecordBatch::try_from_iter(vec![("value", Arc::new(array2) as ArrayRef)]).unwrap();
5829
5830 let mut buffer = Vec::new();
5831 let options = WriterProperties::builder()
5832 .set_max_row_group_row_count(Some(3))
5833 .build();
5834 let mut writer = ArrowWriter::try_new(&mut buffer, batch1.schema(), Some(options)).unwrap();
5835 writer.write(&batch1).unwrap();
5836 writer.write(&batch2).unwrap();
5837 writer.close().unwrap();
5838
5839 let file = Bytes::from(buffer);
5840 let row_group_index_field = Arc::new(
5841 Field::new("row_group_index", ArrowDataType::Int64, false)
5842 .with_extension_type(RowGroupIndex),
5843 );
5844
5845 let options = ArrowReaderOptions::new()
5846 .with_virtual_columns(vec![row_group_index_field.clone()])
5847 .unwrap();
5848 let metadata = ArrowReaderMetadata::load(&file, options).unwrap();
5849 let num_columns = metadata
5850 .metadata
5851 .file_metadata()
5852 .schema_descr()
5853 .num_columns();
5854
5855 let mut arrow_reader = ParquetRecordBatchReaderBuilder::new_with_metadata(file, metadata)
5856 .with_projection(ProjectionMask::none(num_columns))
5857 .build()
5858 .expect("reader with virtual columns only");
5859
5860 let batch = arrow_reader.next().unwrap().unwrap();
5861 let schema = Arc::new(Schema::new(vec![(*row_group_index_field).clone()]));
5862
5863 assert_eq!(batch.schema(), schema);
5864 assert_eq!(batch.num_columns(), 1);
5865 assert_eq!(batch.num_rows(), 5);
5866
5867 assert_eq!(
5868 batch
5869 .column(0)
5870 .as_primitive::<types::Int64Type>()
5871 .iter()
5872 .collect::<Vec<_>>(),
5873 vec![Some(0), Some(0), Some(0), Some(1), Some(1)]
5874 );
5875 }
5876
5877 #[test]
5878 fn test_read_row_group_indices_with_selection() -> Result<()> {
5879 let mut buffer = Vec::new();
5880 let options = WriterProperties::builder()
5881 .set_max_row_group_row_count(Some(10))
5882 .build();
5883
5884 let schema = Arc::new(Schema::new(vec![Field::new(
5885 "value",
5886 ArrowDataType::Int64,
5887 false,
5888 )]));
5889
5890 let mut writer = ArrowWriter::try_new(&mut buffer, schema.clone(), Some(options))?;
5891
5892 for i in 0..3 {
5894 let start = i * 10;
5895 let array = Int64Array::from_iter_values(start..start + 10);
5896 let batch = RecordBatch::try_from_iter(vec![("value", Arc::new(array) as ArrayRef)])?;
5897 writer.write(&batch)?;
5898 }
5899 writer.close()?;
5900
5901 let file = Bytes::from(buffer);
5902 let row_group_index_field = Arc::new(
5903 Field::new("rg_idx", ArrowDataType::Int64, false).with_extension_type(RowGroupIndex),
5904 );
5905
5906 let options =
5907 ArrowReaderOptions::new().with_virtual_columns(vec![row_group_index_field])?;
5908
5909 let arrow_reader =
5911 ParquetRecordBatchReaderBuilder::try_new_with_options(file.clone(), options.clone())?
5912 .with_row_groups(vec![2, 1, 0])
5913 .build()?;
5914
5915 let batches: Vec<_> = arrow_reader.collect::<Result<Vec<_>, _>>()?;
5916 let combined = concat_batches(&batches[0].schema(), &batches)?;
5917
5918 let values = combined.column(0).as_primitive::<types::Int64Type>();
5919 let first_val = values.value(0);
5920 let last_val = values.value(combined.num_rows() - 1);
5921 assert_eq!(first_val, 20);
5923 assert_eq!(last_val, 9);
5925
5926 let rg_indices = combined.column(1).as_primitive::<types::Int64Type>();
5927 assert_eq!(rg_indices.value(0), 2);
5928 assert_eq!(rg_indices.value(10), 1);
5929 assert_eq!(rg_indices.value(20), 0);
5930
5931 Ok(())
5932 }
5933
5934 pub(crate) fn test_row_numbers_with_multiple_row_groups_helper<F>(
5935 use_filter: bool,
5936 test_case: F,
5937 ) where
5938 F: FnOnce(PathBuf, RowSelection, Option<RowFilter>, usize) -> Vec<RecordBatch>,
5939 {
5940 let seed: u64 = random();
5941 println!("test_row_numbers_with_multiple_row_groups seed: {}", seed);
5942 let mut rng = StdRng::seed_from_u64(seed);
5943
5944 use tempfile::TempDir;
5945 let tempdir = TempDir::new().expect("Could not create temp dir");
5946
5947 let (bytes, metadata) = generate_file_with_row_numbers(&mut rng);
5948
5949 let path = tempdir.path().join("test.parquet");
5950 std::fs::write(&path, bytes).expect("Could not write file");
5951
5952 let mut case = vec![];
5953 let mut remaining = metadata.file_metadata().num_rows();
5954 while remaining > 0 {
5955 let row_count = rng.random_range(1..=remaining);
5956 remaining -= row_count;
5957 case.push(RowSelector {
5958 row_count: row_count as usize,
5959 skip: rng.random_bool(0.5),
5960 });
5961 }
5962
5963 let filter = use_filter.then(|| {
5964 let filter = (0..metadata.file_metadata().num_rows())
5965 .map(|_| rng.random_bool(0.99))
5966 .collect::<Vec<_>>();
5967 let mut filter_offset = 0;
5968 RowFilter::new(vec![Box::new(ArrowPredicateFn::new(
5969 ProjectionMask::all(),
5970 move |b| {
5971 let array = BooleanArray::from_iter(
5972 filter
5973 .iter()
5974 .skip(filter_offset)
5975 .take(b.num_rows())
5976 .map(|x| Some(*x)),
5977 );
5978 filter_offset += b.num_rows();
5979 Ok(array)
5980 },
5981 ))])
5982 });
5983
5984 let selection = RowSelection::from(case);
5985 let batches = test_case(path, selection.clone(), filter, rng.random_range(1..4096));
5986
5987 if selection.skipped_row_count() == metadata.file_metadata().num_rows() as usize {
5988 assert!(batches.into_iter().all(|batch| batch.num_rows() == 0));
5989 return;
5990 }
5991 let actual = concat_batches(batches.first().expect("No batches").schema_ref(), &batches)
5992 .expect("Failed to concatenate");
5993 let values = actual
5995 .column(0)
5996 .as_primitive::<types::Int64Type>()
5997 .iter()
5998 .collect::<Vec<_>>();
5999 let row_numbers = actual
6000 .column(1)
6001 .as_primitive::<types::Int64Type>()
6002 .iter()
6003 .collect::<Vec<_>>();
6004 assert_eq!(
6005 row_numbers
6006 .into_iter()
6007 .map(|number| number.map(|number| number + 1))
6008 .collect::<Vec<_>>(),
6009 values
6010 );
6011 }
6012
6013 fn generate_file_with_row_numbers(rng: &mut impl Rng) -> (Bytes, ParquetMetaData) {
6014 let schema = Arc::new(Schema::new(Fields::from(vec![Field::new(
6015 "value",
6016 ArrowDataType::Int64,
6017 false,
6018 )])));
6019
6020 let mut buf = Vec::with_capacity(1024);
6021 let mut writer =
6022 ArrowWriter::try_new(&mut buf, schema.clone(), None).expect("Could not create writer");
6023
6024 let mut values = 1..=rng.random_range(1..4096);
6025 while !values.is_empty() {
6026 let batch_values = values
6027 .by_ref()
6028 .take(rng.random_range(1..4096))
6029 .collect::<Vec<_>>();
6030 let array = Arc::new(Int64Array::from(batch_values)) as ArrayRef;
6031 let batch =
6032 RecordBatch::try_from_iter([("value", array)]).expect("Could not create batch");
6033 writer.write(&batch).expect("Could not write batch");
6034 writer.flush().expect("Could not flush");
6035 }
6036 let metadata = writer.close().expect("Could not close writer");
6037
6038 (Bytes::from(buf), metadata)
6039 }
6040}