Skip to main content

parquet/arrow/arrow_reader/
mod.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Contains reader which reads parquet data into arrow [`RecordBatch`]
19
20use 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;
55// Exposed so integration tests and benchmarks can temporarily override the threshold.
56pub use read_plan::{PredicateOptions, ReadPlan, ReadPlanBuilder};
57
58mod filter;
59pub mod metrics;
60mod read_plan;
61pub(crate) mod selection;
62pub mod statistics;
63
64/// Default batch size for reading parquet files
65pub const DEFAULT_BATCH_SIZE: usize = 1024;
66
67/// Builder for constructing Parquet readers that decode into [Apache Arrow]
68/// arrays.
69///
70/// Most users should use one of the following specializations:
71///
72/// * synchronous API: [`ParquetRecordBatchReaderBuilder`]
73/// * `async` API: [`ParquetRecordBatchStreamBuilder`]
74/// * decoder API: [`ParquetPushDecoderBuilder`]
75///
76/// # Features
77/// * Projection pushdown: [`Self::with_projection`]
78/// * Cached metadata: [`ArrowReaderMetadata::load`]
79/// * Offset skipping: [`Self::with_offset`] and [`Self::with_limit`]
80/// * Row group filtering: [`Self::with_row_groups`]
81/// * Range filtering: [`Self::with_row_selection`]
82/// * Row level filtering: [`Self::with_row_filter`]
83///
84/// # Implementing Predicate Pushdown
85///
86/// [`Self::with_row_filter`] permits filter evaluation *during* the decoding
87/// process, which is efficient and allows the most low level optimizations.
88///
89/// However, most Parquet based systems will apply filters at many steps prior
90/// to decoding such as pruning files, row groups and data pages. This crate
91/// provides the low level APIs needed to implement such filtering, but does not
92/// include any logic to actually evaluate predicates. For example:
93///
94/// * [`Self::with_row_groups`] for Row Group pruning
95/// * [`Self::with_row_selection`] for data page pruning
96/// * [`StatisticsConverter`] to convert Parquet statistics to Arrow arrays
97///
98/// The rationale for this design is that implementing predicate pushdown is a
99/// complex topic and varies significantly from system to system. For example
100///
101/// 1. Predicates supported (do you support predicates like prefix matching, user defined functions, etc)
102/// 2. Evaluating predicates on multiple files (with potentially different but compatible schemas)
103/// 3. Evaluating predicates using information from an external metadata catalog (e.g. Apache Iceberg or similar)
104/// 4. Interleaving fetching metadata, evaluating predicates, and decoding files
105///
106/// You can read more about this design in the [Querying Parquet with
107/// Millisecond Latency] Arrow blog post.
108///
109/// [`ParquetRecordBatchStreamBuilder`]: crate::arrow::async_reader::ParquetRecordBatchStreamBuilder
110/// [`ParquetPushDecoderBuilder`]: crate::arrow::push_decoder::ParquetPushDecoderBuilder
111/// [Apache Arrow]: https://arrow.apache.org/
112/// [`StatisticsConverter`]: statistics::StatisticsConverter
113/// [Querying Parquet with Millisecond Latency]: https://arrow.apache.org/blog/2022/12/26/querying-parquet-with-millisecond-latency/
114pub struct ArrowReaderBuilder<T> {
115    /// The "input" to read parquet data from.
116    ///
117    /// Note in the case of the [`ParquetPushDecoderBuilder`] there is no
118    /// underlying reader; the input is instead [`PushDecoderInput`], the buffer that
119    /// caller-pushed bytes accumulate in.
120    ///
121    /// [`ParquetPushDecoderBuilder`]: crate::arrow::push_decoder::ParquetPushDecoderBuilder
122    /// [`PushDecoderInput`]: crate::arrow::push_decoder::PushDecoderInput
123    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, // 100MB default cache size
189        }
190    }
191
192    /// Returns a reference to the [`ParquetMetaData`] for this parquet file
193    pub fn metadata(&self) -> &Arc<ParquetMetaData> {
194        &self.metadata
195    }
196
197    /// Returns the parquet [`SchemaDescriptor`] for this parquet file
198    pub fn parquet_schema(&self) -> &SchemaDescriptor {
199        self.metadata.file_metadata().schema_descr()
200    }
201
202    /// Returns the arrow [`SchemaRef`] for this parquet file
203    pub fn schema(&self) -> &SchemaRef {
204        &self.schema
205    }
206
207    /// Set the size of [`RecordBatch`] to produce. Defaults to [`DEFAULT_BATCH_SIZE`].
208    ///
209    /// This may be used as a hint for internal allocations, but does not
210    /// guarantee exact internal buffer capacities.
211    ///
212    /// If `batch_size` is more than the file row count, use the file row count.
213    pub fn with_batch_size(self, batch_size: usize) -> Self {
214        // Try to avoid allocate large buffer
215        let batch_size = batch_size.min(self.metadata.file_metadata().num_rows() as usize);
216        Self { batch_size, ..self }
217    }
218
219    /// Only read data from the provided row group indexes
220    ///
221    /// This is also called row group filtering
222    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    /// Only read data from the provided column indexes
230    pub fn with_projection(self, mask: ProjectionMask) -> Self {
231        Self {
232            projection: mask,
233            ..self
234        }
235    }
236
237    /// Configure how row selections should be materialised during execution
238    ///
239    /// See [`RowSelectionPolicy`] for more details
240    pub fn with_row_selection_policy(self, policy: RowSelectionPolicy) -> Self {
241        Self {
242            row_selection_policy: policy,
243            ..self
244        }
245    }
246
247    /// Provide a [`RowSelection`] to filter out rows, and avoid fetching their
248    /// data into memory.
249    ///
250    /// This feature is used to restrict which rows are decoded within row
251    /// groups, skipping ranges of rows that are not needed. Such selections
252    /// could be determined by evaluating predicates against the parquet page
253    /// [`Index`] or some other external information available to a query
254    /// engine.
255    ///
256    /// # Notes
257    ///
258    /// Row group filtering (see [`Self::with_row_groups`]) is applied prior to
259    /// applying the row selection, and therefore rows from skipped row groups
260    /// should not be included in the [`RowSelection`] (see example below)
261    ///
262    /// It is recommended to enable writing the page index if using this
263    /// functionality, to allow more efficient skipping over data pages. See
264    /// [`ArrowReaderOptions::with_page_index`].
265    ///
266    /// # Example
267    ///
268    /// Given a parquet file with 4 row groups, and a row group filter of `[0,
269    /// 2, 3]`, in order to scan rows 50-100 in row group 2 and rows 200-300 in
270    /// row group 3:
271    ///
272    /// ```text
273    ///   Row Group 0, 1000 rows (selected)
274    ///   Row Group 1, 1000 rows (skipped)
275    ///   Row Group 2, 1000 rows (selected, but want to only scan rows 50-100)
276    ///   Row Group 3, 1000 rows (selected, but want to only scan rows 200-300)
277    /// ```
278    ///
279    /// You could pass the following [`RowSelection`]:
280    ///
281    /// ```text
282    ///  Select 1000    (scan all rows in row group 0)
283    ///  Skip 50        (skip the first 50 rows in row group 2)
284    ///  Select 50      (scan rows 50-100 in row group 2)
285    ///  Skip 900       (skip the remaining rows in row group 2)
286    ///  Skip 200       (skip the first 200 rows in row group 3)
287    ///  Select 100     (scan rows 200-300 in row group 3)
288    ///  Skip 700       (skip the remaining rows in row group 3)
289    /// ```
290    /// Note there is no entry for the (entirely) skipped row group 1.
291    ///
292    /// Note you can represent the same selection with fewer entries. Instead of
293    ///
294    /// ```text
295    ///  Skip 900       (skip the remaining rows in row group 2)
296    ///  Skip 200       (skip the first 200 rows in row group 3)
297    /// ```
298    ///
299    /// you could use
300    ///
301    /// ```text
302    /// Skip 1100      (skip the remaining 900 rows in row group 2 and the first 200 rows in row group 3)
303    /// ```
304    ///
305    /// [`Index`]: crate::file::page_index::column_index::ColumnIndexMetaData
306    pub fn with_row_selection(self, selection: RowSelection) -> Self {
307        Self {
308            selection: Some(selection),
309            ..self
310        }
311    }
312
313    /// Provide a [`RowFilter`] to skip decoding rows
314    ///
315    /// Row filters are applied after row group selection and row selection
316    ///
317    /// It is recommended to enable reading the page index if using this functionality, to allow
318    /// more efficient skipping over data pages. See [`ArrowReaderOptions::with_page_index`].
319    ///
320    /// See the [blog post on late materialization] for a more technical explanation.
321    ///
322    /// [blog post on late materialization]: https://arrow.apache.org/blog/2025/12/11/parquet-late-materialization-deep-dive
323    ///
324    /// # Example
325    /// ```rust
326    /// # use std::fs::File;
327    /// # use arrow_array::Int32Array;
328    /// # use parquet::arrow::ProjectionMask;
329    /// # use parquet::arrow::arrow_reader::{ArrowPredicateFn, ParquetRecordBatchReaderBuilder, RowFilter};
330    /// # fn main() -> Result<(), parquet::errors::ParquetError> {
331    /// # let testdata = arrow::util::test_util::parquet_test_data();
332    /// # let path = format!("{testdata}/alltypes_plain.parquet");
333    /// # let file = File::open(&path)?;
334    /// let builder = ParquetRecordBatchReaderBuilder::try_new(file)?;
335    /// let schema_desc = builder.metadata().file_metadata().schema_descr_ptr();
336    /// // Create predicate that evaluates `int_col != 1`.
337    /// // `int_col` column has index 4 (zero based) in the schema
338    /// let projection = ProjectionMask::leaves(&schema_desc, [4]);
339    /// // Only the projection columns are passed to the predicate so
340    /// // int_col is column 0 in the predicate
341    /// let predicate = ArrowPredicateFn::new(projection, |batch| {
342    ///     let int_col = batch.column(0);
343    ///     arrow::compute::kernels::cmp::neq(int_col, &Int32Array::new_scalar(1))
344    /// });
345    /// let row_filter = RowFilter::new(vec![Box::new(predicate)]);
346    /// // The filter will be invoked during the reading process
347    /// let reader = builder.with_row_filter(row_filter).build()?;
348    /// # for b in reader { let _ = b?; }
349    /// # Ok(())
350    /// # }
351    /// ```
352    pub fn with_row_filter(self, filter: RowFilter) -> Self {
353        Self {
354            filter: Some(filter),
355            ..self
356        }
357    }
358
359    /// Provide a limit to the number of rows to be read
360    ///
361    /// The limit will be applied after any [`Self::with_row_selection`] and [`Self::with_row_filter`]
362    /// allowing it to limit the final set of rows decoded after any pushed down predicates
363    ///
364    /// It is recommended to enable reading the page index if using this functionality, to allow
365    /// more efficient skipping over data pages. See [`ArrowReaderOptions::with_page_index`]
366    pub fn with_limit(self, limit: usize) -> Self {
367        Self {
368            limit: Some(limit),
369            ..self
370        }
371    }
372
373    /// Provide an offset to skip over the given number of rows
374    ///
375    /// The offset will be applied after any [`Self::with_row_selection`] and [`Self::with_row_filter`]
376    /// allowing it to skip rows after any pushed down predicates
377    ///
378    /// It is recommended to enable reading the page index if using this functionality, to allow
379    /// more efficient skipping over data pages. See [`ArrowReaderOptions::with_page_index`]
380    pub fn with_offset(self, offset: usize) -> Self {
381        Self {
382            offset: Some(offset),
383            ..self
384        }
385    }
386
387    /// Specify metrics collection during reading
388    ///
389    /// To access the metrics, create an [`ArrowReaderMetrics`] and pass a
390    /// clone of the provided metrics to the builder.
391    ///
392    /// For example:
393    ///
394    /// ```rust
395    /// # use std::sync::Arc;
396    /// # use bytes::Bytes;
397    /// # use arrow_array::{Int32Array, RecordBatch};
398    /// # use arrow_schema::{DataType, Field, Schema};
399    /// # use parquet::arrow::arrow_reader::{ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder};
400    /// use parquet::arrow::arrow_reader::metrics::ArrowReaderMetrics;
401    /// # use parquet::arrow::ArrowWriter;
402    /// # let mut file: Vec<u8> = Vec::with_capacity(1024);
403    /// # let schema = Arc::new(Schema::new(vec![Field::new("i32", DataType::Int32, false)]));
404    /// # let mut writer = ArrowWriter::try_new(&mut file, schema.clone(), None).unwrap();
405    /// # let batch = RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1, 2, 3]))]).unwrap();
406    /// # writer.write(&batch).unwrap();
407    /// # writer.close().unwrap();
408    /// # let file = Bytes::from(file);
409    /// // Create metrics object to pass into the reader
410    /// let metrics = ArrowReaderMetrics::enabled();
411    /// let reader = ParquetRecordBatchReaderBuilder::try_new(file).unwrap()
412    ///   // Configure the builder to use the metrics by passing a clone
413    ///   .with_metrics(metrics.clone())
414    ///   // Build the reader
415    ///   .build().unwrap();
416    /// // .. read data from the reader ..
417    ///
418    /// // check the metrics
419    /// assert!(metrics.records_read_from_inner().is_some());
420    /// ```
421    pub fn with_metrics(self, metrics: ArrowReaderMetrics) -> Self {
422        Self { metrics, ..self }
423    }
424
425    /// Set the maximum size (per row group) of the predicate cache in bytes for
426    /// the async decoder.
427    ///
428    /// Defaults to 100MB (across all columns). Set to `usize::MAX` to use
429    /// unlimited cache size.
430    ///
431    /// This cache is used to store decoded arrays that are used in
432    /// predicate evaluation ([`Self::with_row_filter`]).
433    ///
434    /// This cache is only used for the "async" decoder, [`ParquetRecordBatchStream`]. See
435    /// [this ticket] for more details and alternatives.
436    ///
437    /// [`ParquetRecordBatchStream`]: https://docs.rs/parquet/latest/parquet/arrow/async_reader/struct.ParquetRecordBatchStream.html
438    /// [this ticket]: https://github.com/apache/arrow-rs/issues/8000
439    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/// Options that control how [`ParquetMetaData`] is read when constructing
448/// an Arrow reader.
449///
450/// To use these options, pass them to one of the following methods:
451/// * [`ParquetRecordBatchReaderBuilder::try_new_with_options`]
452/// * [`ParquetRecordBatchStreamBuilder::new_with_options`]
453///
454/// For fine-grained control over metadata loading, use
455/// [`ArrowReaderMetadata::load`] to load metadata with these options,
456///
457/// See [`ArrowReaderBuilder`] for how to configure how the column data
458/// is then read from the file, including projection and filter pushdown
459///
460/// [`ParquetRecordBatchStreamBuilder::new_with_options`]: crate::arrow::async_reader::ParquetRecordBatchStreamBuilder::new_with_options
461#[derive(Debug, Clone, Default)]
462pub struct ArrowReaderOptions {
463    /// Should the reader strip any user defined metadata from the Arrow schema
464    skip_arrow_metadata: bool,
465    /// If provided, used as the schema hint when determining the Arrow schema,
466    /// otherwise the schema hint is read from the [ARROW_SCHEMA_META_KEY]
467    ///
468    /// [ARROW_SCHEMA_META_KEY]: crate::arrow::ARROW_SCHEMA_META_KEY
469    supplied_schema: Option<SchemaRef>,
470
471    pub(crate) column_index: PageIndexPolicy,
472    pub(crate) offset_index: PageIndexPolicy,
473
474    /// Options to control reading of Parquet metadata
475    metadata_options: ParquetMetaDataOptions,
476    /// If encryption is enabled, the file decryption properties can be provided
477    #[cfg(feature = "encryption")]
478    pub(crate) file_decryption_properties: Option<Arc<FileDecryptionProperties>>,
479
480    virtual_columns: Vec<FieldRef>,
481}
482
483impl ArrowReaderOptions {
484    /// Create a new [`ArrowReaderOptions`] with the default settings
485    pub fn new() -> Self {
486        Self::default()
487    }
488
489    /// Skip decoding the embedded arrow metadata (defaults to `false`)
490    ///
491    /// Parquet files generated by some writers may contain embedded arrow
492    /// schema and metadata.
493    /// This may not be correct or compatible with your system,
494    /// for example, see [ARROW-16184](https://issues.apache.org/jira/browse/ARROW-16184)
495    pub fn with_skip_arrow_metadata(self, skip_arrow_metadata: bool) -> Self {
496        Self {
497            skip_arrow_metadata,
498            ..self
499        }
500    }
501
502    /// Provide a schema hint to use when reading the Parquet file.
503    ///
504    /// If provided, this schema takes precedence over any arrow schema embedded
505    /// in the metadata (see the [`arrow`] documentation for more details).
506    ///
507    /// If the provided schema is not compatible with the data stored in the
508    /// parquet file schema, an error will be returned when constructing the
509    /// builder.
510    ///
511    /// This option is only required if you want to explicitly control the
512    /// conversion of Parquet types to Arrow types, such as casting a column to
513    /// a different type. For example, if you wanted to read an Int64 in
514    /// a Parquet file to a [`TimestampMicrosecondArray`] in the Arrow schema.
515    ///
516    /// [`arrow`]: crate::arrow
517    /// [`TimestampMicrosecondArray`]: arrow_array::TimestampMicrosecondArray
518    ///
519    /// # Notes
520    ///
521    /// The provided schema must have the same number of columns as the parquet schema and
522    /// the column names must be the same.
523    ///
524    /// # Example
525    /// ```
526    /// # use std::sync::Arc;
527    /// # use bytes::Bytes;
528    /// # use arrow_array::{ArrayRef, Int32Array, RecordBatch};
529    /// # use arrow_schema::{DataType, Field, Schema, TimeUnit};
530    /// # use parquet::arrow::arrow_reader::{ArrowReaderOptions, ParquetRecordBatchReaderBuilder};
531    /// # use parquet::arrow::ArrowWriter;
532    /// // Write data - schema is inferred from the data to be Int32
533    /// let mut file = Vec::new();
534    /// let batch = RecordBatch::try_from_iter(vec![
535    ///     ("col_1", Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef),
536    /// ]).unwrap();
537    /// let mut writer = ArrowWriter::try_new(&mut file, batch.schema(), None).unwrap();
538    /// writer.write(&batch).unwrap();
539    /// writer.close().unwrap();
540    /// let file = Bytes::from(file);
541    ///
542    /// // Read the file back.
543    /// // Supply a schema that interprets the Int32 column as a Timestamp.
544    /// let supplied_schema = Arc::new(Schema::new(vec![
545    ///     Field::new("col_1", DataType::Timestamp(TimeUnit::Nanosecond, None), false)
546    /// ]));
547    /// let options = ArrowReaderOptions::new().with_schema(supplied_schema.clone());
548    /// let mut builder = ParquetRecordBatchReaderBuilder::try_new_with_options(
549    ///     file.clone(),
550    ///     options
551    /// ).expect("Error if the schema is not compatible with the parquet file schema.");
552    ///
553    /// // Create the reader and read the data using the supplied schema.
554    /// let mut reader = builder.build().unwrap();
555    /// let _batch = reader.next().unwrap().unwrap();
556    /// ```
557    ///
558    /// # Example: Preserving Dictionary Encoding
559    ///
560    /// By default, Parquet string columns are read as `Utf8Array` (or `LargeUtf8Array`),
561    /// even if the underlying Parquet data uses dictionary encoding. You can preserve
562    /// the dictionary encoding by specifying a `Dictionary` type in the schema hint:
563    ///
564    /// ```
565    /// # use std::sync::Arc;
566    /// # use bytes::Bytes;
567    /// # use arrow_array::{ArrayRef, RecordBatch, StringArray};
568    /// # use arrow_schema::{DataType, Field, Schema};
569    /// # use parquet::arrow::arrow_reader::{ArrowReaderOptions, ParquetRecordBatchReaderBuilder};
570    /// # use parquet::arrow::ArrowWriter;
571    /// // Write a Parquet file with string data
572    /// let mut file = Vec::new();
573    /// let schema = Arc::new(Schema::new(vec![
574    ///     Field::new("city", DataType::Utf8, false)
575    /// ]));
576    /// let cities = StringArray::from(vec!["Berlin", "Berlin", "Paris", "Berlin", "Paris"]);
577    /// let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(cities)]).unwrap();
578    ///
579    /// let mut writer = ArrowWriter::try_new(&mut file, batch.schema(), None).unwrap();
580    /// writer.write(&batch).unwrap();
581    /// writer.close().unwrap();
582    /// let file = Bytes::from(file);
583    ///
584    /// // Read the file back, requesting dictionary encoding preservation
585    /// let dict_schema = Arc::new(Schema::new(vec![
586    ///     Field::new("city", DataType::Dictionary(
587    ///         Box::new(DataType::Int32),
588    ///         Box::new(DataType::Utf8)
589    ///     ), false)
590    /// ]));
591    /// let options = ArrowReaderOptions::new().with_schema(dict_schema);
592    /// let builder = ParquetRecordBatchReaderBuilder::try_new_with_options(
593    ///     file.clone(),
594    ///     options
595    /// ).unwrap();
596    ///
597    /// let mut reader = builder.build().unwrap();
598    /// let batch = reader.next().unwrap().unwrap();
599    ///
600    /// // The column is now a DictionaryArray
601    /// assert!(matches!(
602    ///     batch.column(0).data_type(),
603    ///     DataType::Dictionary(_, _)
604    /// ));
605    /// ```
606    ///
607    /// **Note**: Dictionary encoding preservation works best when:
608    /// 1. The original column was dictionary encoded (the default for string columns)
609    /// 2. There are a small number of distinct values
610    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    /// Enable reading the [`PageIndex`] from the metadata, if present (defaults to `false`)
620    ///
621    /// The `PageIndex` can be used to push down predicates to the parquet scan,
622    /// potentially eliminating unnecessary IO, by some query engines.
623    ///
624    /// If this is enabled, [`ParquetMetaData::column_index`] and
625    /// [`ParquetMetaData::offset_index`] will be populated if the corresponding
626    /// information is present in the file.
627    ///
628    /// [`PageIndex`]: https://github.com/apache/parquet-format/blob/master/PageIndex.md
629    /// [`ParquetMetaData::column_index`]: crate::file::metadata::ParquetMetaData::column_index
630    /// [`ParquetMetaData::offset_index`]: crate::file::metadata::ParquetMetaData::offset_index
631    pub fn with_page_index(self, page_index: bool) -> Self {
632        self.with_page_index_policy(PageIndexPolicy::from(page_index))
633    }
634
635    /// Sets the [`PageIndexPolicy`] for both the column and offset indexes.
636    ///
637    /// The `PageIndex` consists of two structures: the `ColumnIndex` and `OffsetIndex`.
638    /// This method sets the same policy for both. For fine-grained control, use
639    /// [`Self::with_column_index_policy`] and [`Self::with_offset_index_policy`].
640    ///
641    /// See [`Self::with_page_index`] for more details on page indexes.
642    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    /// Sets the [`PageIndexPolicy`] for the Parquet [ColumnIndex] structure.
648    ///
649    /// The `ColumnIndex` contains min/max statistics for each page, which can be used
650    /// for predicate pushdown and page-level pruning.
651    ///
652    /// [ColumnIndex]: https://github.com/apache/parquet-format/blob/master/PageIndex.md
653    pub fn with_column_index_policy(mut self, policy: PageIndexPolicy) -> Self {
654        self.column_index = policy;
655        self
656    }
657
658    /// Sets the [`PageIndexPolicy`] for the Parquet [OffsetIndex] structure.
659    ///
660    /// The `OffsetIndex` contains the locations and sizes of each page, which enables
661    /// efficient page-level skipping and random access within column chunks.
662    ///
663    /// [OffsetIndex]: https://github.com/apache/parquet-format/blob/master/PageIndex.md
664    pub fn with_offset_index_policy(mut self, policy: PageIndexPolicy) -> Self {
665        self.offset_index = policy;
666        self
667    }
668
669    /// Provide a Parquet schema to use when decoding the metadata. The schema in the Parquet
670    /// footer will be skipped.
671    ///
672    /// This can be used to avoid reparsing the schema from the file when it is
673    /// already known.
674    pub fn with_parquet_schema(mut self, schema: Arc<SchemaDescriptor>) -> Self {
675        self.metadata_options.set_schema(schema);
676        self
677    }
678
679    /// Set whether to convert the [`encoding_stats`] in the Parquet `ColumnMetaData` to a bitmask
680    /// (defaults to `false`).
681    ///
682    /// See [`ColumnChunkMetaData::page_encoding_stats_mask`] for an explanation of why this
683    /// might be desirable.
684    ///
685    /// [`ColumnChunkMetaData::page_encoding_stats_mask`]:
686    /// crate::file::metadata::ColumnChunkMetaData::page_encoding_stats_mask
687    /// [`encoding_stats`]:
688    /// https://github.com/apache/parquet-format/blob/786142e26740487930ddc3ec5e39d780bd930907/src/main/thrift/parquet.thrift#L917
689    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    /// Sets the decoding policy for [`encoding_stats`] in the Parquet `ColumnMetaData`.
695    ///
696    /// [`encoding_stats`]:
697    /// https://github.com/apache/parquet-format/blob/786142e26740487930ddc3ec5e39d780bd930907/src/main/thrift/parquet.thrift#L917
698    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    /// Sets the decoding policy for [`statistics`] in the Parquet `ColumnMetaData`.
704    ///
705    /// [`statistics`]:
706    /// https://github.com/apache/parquet-format/blob/786142e26740487930ddc3ec5e39d780bd930907/src/main/thrift/parquet.thrift#L912
707    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    /// Sets the decoding policy for [`size_statistics`] in the Parquet `ColumnMetaData`.
713    ///
714    /// [`size_statistics`]:
715    /// https://github.com/apache/parquet-format/blob/786142e26740487930ddc3ec5e39d780bd930907/src/main/thrift/parquet.thrift#L936
716    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    /// Provide the file decryption properties to use when reading encrypted parquet files.
722    ///
723    /// If encryption is enabled and the file is encrypted, the `file_decryption_properties` must be provided.
724    #[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    /// Include virtual columns in the output.
736    ///
737    /// Virtual columns are columns that are not part of the Parquet schema, but are added to the output by the reader such as row numbers and row group indices.
738    ///
739    /// # Example
740    /// ```
741    /// # use std::sync::Arc;
742    /// # use bytes::Bytes;
743    /// # use arrow_array::{ArrayRef, Int64Array, RecordBatch};
744    /// # use arrow_schema::{DataType, Field, Schema};
745    /// # use parquet::arrow::{ArrowWriter, RowNumber};
746    /// # use parquet::arrow::arrow_reader::{ArrowReaderOptions, ParquetRecordBatchReaderBuilder};
747    /// #
748    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
749    /// // Create a simple record batch with some data
750    /// let values = Arc::new(Int64Array::from(vec![1, 2, 3])) as ArrayRef;
751    /// let batch = RecordBatch::try_from_iter(vec![("value", values)])?;
752    ///
753    /// // Write the batch to an in-memory buffer
754    /// let mut file = Vec::new();
755    /// let mut writer = ArrowWriter::try_new(
756    ///     &mut file,
757    ///     batch.schema(),
758    ///     None
759    /// )?;
760    /// writer.write(&batch)?;
761    /// writer.close()?;
762    /// let file = Bytes::from(file);
763    ///
764    /// // Create a virtual column for row numbers
765    /// let row_number_field = Arc::new(Field::new("row_number", DataType::Int64, false)
766    ///     .with_extension_type(RowNumber));
767    ///
768    /// // Configure options with virtual columns
769    /// let options = ArrowReaderOptions::new()
770    ///     .with_virtual_columns(vec![row_number_field])?;
771    ///
772    /// // Create a reader with the options
773    /// let mut reader = ParquetRecordBatchReaderBuilder::try_new_with_options(
774    ///     file,
775    ///     options
776    /// )?
777    /// .build()?;
778    ///
779    /// // Read the batch - it will include both the original column and the virtual row_number column
780    /// let result_batch = reader.next().unwrap()?;
781    /// assert_eq!(result_batch.num_columns(), 2); // "value" + "row_number"
782    /// assert_eq!(result_batch.num_rows(), 3);
783    /// #
784    /// # Ok(())
785    /// # }
786    /// ```
787    pub fn with_virtual_columns(self, virtual_columns: Vec<FieldRef>) -> Result<Self> {
788        // Validate that all fields are virtual columns
789        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    /// Returns whether page index reading is enabled.
808    ///
809    /// This returns `true` if both the column index and offset index policies are not [`PageIndexPolicy::Skip`].
810    ///
811    /// This can be set via [`with_page_index`][Self::with_page_index] or
812    /// [`with_page_index_policy`][Self::with_page_index_policy].
813    pub fn page_index(&self) -> bool {
814        self.offset_index != PageIndexPolicy::Skip && self.column_index != PageIndexPolicy::Skip
815    }
816
817    /// Retrieve the currently set [`PageIndexPolicy`] for the offset index.
818    ///
819    /// This can be set via [`with_offset_index_policy`][Self::with_offset_index_policy]
820    /// or [`with_page_index_policy`][Self::with_page_index_policy].
821    pub fn offset_index_policy(&self) -> PageIndexPolicy {
822        self.offset_index
823    }
824
825    /// Retrieve the currently set [`PageIndexPolicy`] for the column index.
826    ///
827    /// This can be set via [`with_column_index_policy`][Self::with_column_index_policy]
828    /// or [`with_page_index_policy`][Self::with_page_index_policy].
829    pub fn column_index_policy(&self) -> PageIndexPolicy {
830        self.column_index
831    }
832
833    /// Retrieve the currently set metadata decoding options.
834    pub fn metadata_options(&self) -> &ParquetMetaDataOptions {
835        &self.metadata_options
836    }
837
838    /// Retrieve the currently set file decryption properties.
839    ///
840    /// This can be set via
841    /// [`file_decryption_properties`][Self::with_file_decryption_properties].
842    #[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    /// Applies the metadata related settings from [`ArrowReaderOptions`],
850    /// such as the [`ParquetMetaDataOptions`], decryption properties, and
851    /// [`PageIndexPolicy`] to this reader.
852    ///
853    /// The page index policies are only applied if at least one of them is not
854    /// [`PageIndexPolicy::Skip`], so policies previously configured on this
855    /// reader (e.g. from a preload setting) are preserved when the options do
856    /// not request the page index.
857    ///
858    /// This encodes the canonical way to construct a `ParquetMetaDataReader`
859    /// inside `AsyncFileReader::get_metadata` (available with the `async`
860    /// feature), so implementations outside this crate do not need to
861    /// duplicate it.
862    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/// The metadata necessary to construct a [`ArrowReaderBuilder`]
887///
888/// Note this structure is cheaply clone-able as it consists of several arcs.
889///
890/// This structure allows
891///
892/// 1. Loading metadata for a file once and then using that same metadata to
893///    construct multiple separate readers, for example, to distribute readers
894///    across multiple threads
895///
896/// 2. Using a cached copy of the [`ParquetMetadata`] rather than reading it
897///    from the file each time a reader is constructed.
898///
899/// [`ParquetMetadata`]: crate::file::metadata::ParquetMetaData
900#[derive(Debug, Clone)]
901pub struct ArrowReaderMetadata {
902    /// The Parquet Metadata, if known aprior
903    pub(crate) metadata: Arc<ParquetMetaData>,
904    /// The Arrow Schema
905    pub(crate) schema: SchemaRef,
906    /// The Parquet schema (root field)
907    pub(crate) fields: Option<Arc<ParquetField>>,
908}
909
910impl ArrowReaderMetadata {
911    /// Create [`ArrowReaderMetadata`] from the provided [`ArrowReaderOptions`]
912    /// and [`ChunkReader`]
913    ///
914    /// See [`ParquetRecordBatchReaderBuilder::new_with_metadata`] for an
915    /// example of how this can be used
916    ///
917    /// # Notes
918    ///
919    /// If `options` has [`ArrowReaderOptions::with_page_index`] true, but
920    /// `Self::metadata` is missing the page index, this function will attempt
921    /// to load the page index by making an object store request.
922    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    /// Create a new [`ArrowReaderMetadata`] from a pre-existing
936    /// [`ParquetMetaData`] and [`ArrowReaderOptions`].
937    ///
938    /// # Notes
939    ///
940    /// This function will not attempt to load the PageIndex if not present in the metadata, regardless
941    /// of the settings in `options`. See [`Self::load`] to load metadata including the page index if needed.
942    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        // Ensure the supplied schema has the same number of columns as the parquet schema.
987        // parquet_to_arrow_field_levels is expected to throw an error if the schemas have
988        // different lengths, but we check here to be safe.
989        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    /// Returns a reference to the [`ParquetMetaData`] for this parquet file
1042    pub fn metadata(&self) -> &Arc<ParquetMetaData> {
1043        &self.metadata
1044    }
1045
1046    /// Returns the parquet [`SchemaDescriptor`] for this parquet file
1047    pub fn parquet_schema(&self) -> &SchemaDescriptor {
1048        self.metadata.file_metadata().schema_descr()
1049    }
1050
1051    /// Returns the arrow [`SchemaRef`] for this parquet file
1052    pub fn schema(&self) -> &SchemaRef {
1053        &self.schema
1054    }
1055}
1056
1057#[doc(hidden)]
1058// A newtype used within `ReaderOptionsBuilder` to distinguish sync readers from async
1059pub 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
1067/// Creates [`ParquetRecordBatchReader`] for reading Parquet files into Arrow [`RecordBatch`]es
1068///
1069/// # See Also
1070/// * [`crate::arrow::async_reader::ParquetRecordBatchStreamBuilder`] for an async API
1071/// * [`crate::arrow::push_decoder::ParquetPushDecoderBuilder`] for a SansIO decoder API
1072/// * [`ArrowReaderBuilder`] for additional member functions
1073pub type ParquetRecordBatchReaderBuilder<T> = ArrowReaderBuilder<SyncReader<T>>;
1074
1075impl<T: ChunkReader + 'static> ParquetRecordBatchReaderBuilder<T> {
1076    /// Create a new [`ParquetRecordBatchReaderBuilder`]
1077    ///
1078    /// ```
1079    /// # use std::sync::Arc;
1080    /// # use bytes::Bytes;
1081    /// # use arrow_array::{Int32Array, RecordBatch};
1082    /// # use arrow_schema::{DataType, Field, Schema};
1083    /// # use parquet::arrow::arrow_reader::{ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder};
1084    /// # use parquet::arrow::ArrowWriter;
1085    /// # let mut file: Vec<u8> = Vec::with_capacity(1024);
1086    /// # let schema = Arc::new(Schema::new(vec![Field::new("i32", DataType::Int32, false)]));
1087    /// # let mut writer = ArrowWriter::try_new(&mut file, schema.clone(), None).unwrap();
1088    /// # let batch = RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1, 2, 3]))]).unwrap();
1089    /// # writer.write(&batch).unwrap();
1090    /// # writer.close().unwrap();
1091    /// # let file = Bytes::from(file);
1092    /// // Build the reader from anything that implements `ChunkReader`
1093    /// // such as a `File`, or `Bytes`
1094    /// let mut builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
1095    /// // The builder has access to ParquetMetaData such
1096    /// // as the number and layout of row groups
1097    /// assert_eq!(builder.metadata().num_row_groups(), 1);
1098    /// // Call build to create the reader
1099    /// let mut reader: ParquetRecordBatchReader = builder.build().unwrap();
1100    /// // Read data
1101    /// while let Some(batch) = reader.next().transpose()? {
1102    ///     println!("Read {} rows", batch.num_rows());
1103    /// }
1104    /// # Ok::<(), parquet::errors::ParquetError>(())
1105    /// ```
1106    pub fn try_new(reader: T) -> Result<Self> {
1107        Self::try_new_with_options(reader, Default::default())
1108    }
1109
1110    /// Create a new [`ParquetRecordBatchReaderBuilder`] with [`ArrowReaderOptions`]
1111    ///
1112    /// Use this method if you want to control the options for reading the
1113    /// [`ParquetMetaData`]
1114    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    /// Create a [`ParquetRecordBatchReaderBuilder`] from the provided [`ArrowReaderMetadata`]
1120    ///
1121    /// Use this method if you already have [`ParquetMetaData`] for a file.
1122    /// This interface allows:
1123    ///
1124    /// 1. Loading metadata once and using it to create multiple builders with
1125    ///    potentially different settings or run on different threads
1126    ///
1127    /// 2. Using a cached copy of the metadata rather than re-reading it from the
1128    ///    file each time a reader is constructed.
1129    ///
1130    /// See the docs on [`ArrowReaderMetadata`] for more details
1131    ///
1132    /// # Example
1133    /// ```
1134    /// # use std::fs::metadata;
1135    /// # use std::sync::Arc;
1136    /// # use bytes::Bytes;
1137    /// # use arrow_array::{Int32Array, RecordBatch};
1138    /// # use arrow_schema::{DataType, Field, Schema};
1139    /// # use parquet::arrow::arrow_reader::{ArrowReaderMetadata, ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder};
1140    /// # use parquet::arrow::ArrowWriter;
1141    /// #
1142    /// # let mut file: Vec<u8> = Vec::with_capacity(1024);
1143    /// # let schema = Arc::new(Schema::new(vec![Field::new("i32", DataType::Int32, false)]));
1144    /// # let mut writer = ArrowWriter::try_new(&mut file, schema.clone(), None).unwrap();
1145    /// # let batch = RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1, 2, 3]))]).unwrap();
1146    /// # writer.write(&batch).unwrap();
1147    /// # writer.close().unwrap();
1148    /// # let file = Bytes::from(file);
1149    /// #
1150    /// let metadata = ArrowReaderMetadata::load(&file, Default::default()).unwrap();
1151    /// let mut a = ParquetRecordBatchReaderBuilder::new_with_metadata(file.clone(), metadata.clone()).build().unwrap();
1152    /// let mut b = ParquetRecordBatchReaderBuilder::new_with_metadata(file, metadata).build().unwrap();
1153    ///
1154    /// // Should be able to read from both in parallel
1155    /// assert_eq!(a.next().unwrap().unwrap(), b.next().unwrap().unwrap());
1156    /// ```
1157    pub fn new_with_metadata(input: T, metadata: ArrowReaderMetadata) -> Self {
1158        Self::new_builder(SyncReader(input), metadata)
1159    }
1160
1161    /// Read bloom filter for a column in a row group
1162    ///
1163    /// Returns `None` if the column does not have a bloom filter
1164    ///
1165    /// We should call this function after other forms pruning, such as projection and predicate pushdown.
1166    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                // this match exists to future proof the singleton algorithm enum
1193            }
1194        }
1195        match header.compression {
1196            BloomFilterCompression::UNCOMPRESSED => {
1197                // this match exists to future proof the singleton compression enum
1198            }
1199        }
1200        match header.hash {
1201            BloomFilterHash::XXHASH => {
1202                // this match exists to future proof the singleton hash enum
1203            }
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    /// Build a [`ParquetRecordBatchReader`]
1222    ///
1223    /// Note: this will eagerly evaluate any `RowFilter` before returning
1224    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            // Not used for the sync reader, see https://github.com/apache/arrow-rs/issues/8000
1240            max_predicate_cache_size: _,
1241        } = self;
1242
1243        // Try to avoid allocate large buffer
1244        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        // Update selection based on any filters
1259        if let Some(filter) = filter.as_mut() {
1260            for predicate in filter.predicates.iter_mut() {
1261                // break early if we have ruled out all rows
1262                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    /// Optional list of row group indices to scan
1299    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    /// Return the next SerializedPageReader
1342    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        // `offset_index` may not exist and `i[rg_idx]` will be empty.
1347        // To avoid `i[rg_idx][self.column_idx`] panic, we need to filter out empty `i[rg_idx]`.
1348        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
1378/// Reads Parquet data as Arrow [`RecordBatch`]es
1379///
1380/// This struct implements the [`RecordBatchReader`] trait and is an
1381/// `Iterator<Item = ArrowResult<RecordBatch>>` that yields [`RecordBatch`]es.
1382///
1383/// Typically, either reads from a file or an in memory buffer [`Bytes`]
1384///
1385/// Created by [`ParquetRecordBatchReaderBuilder`]
1386///
1387/// [`Bytes`]: bytes::Bytes
1388pub struct ParquetRecordBatchReader {
1389    array_reader: Box<dyn ArrayReader>,
1390    schema: SchemaRef,
1391    read_plan: ReadPlan,
1392}
1393
1394/// Accumulates filter masks for decoded chunks in one logical output batch.
1395///
1396/// The first chunk keeps its [`BooleanBuffer`] without copying. A second chunk
1397/// promotes the accumulator to a [`BooleanBufferBuilder`], and later chunks are
1398/// appended to it. For example, chunks `1000` and `1` become `10001`:
1399///
1400/// ```text
1401///           append(1000)             append(1)
1402///   Empty ───────────────▶ Single ───────────────▶ Combined
1403///                          1000                    10001
1404///                          (zero copy)             (promoted to builder)
1405/// ```
1406///
1407/// The combined mask lines up with the rows that [`ArrayReader::read_records`]
1408/// buffered across the decoded chunks (see [`read_mask_batch`]). Consuming the
1409/// buffered batch and filtering it with the accumulated mask yields the output.
1410///
1411/// ```text
1412///   decoded rows:   0 1 2 3   11   <-- buffered by the array reader
1413///   chunk masks:   [1 0 0 0] [1]
1414///   finish():       1 0 0 0   1    <-- filters the whole batch in one pass
1415/// ```
1416#[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
1450/// Converts the projection buffered by `array_reader` into a record batch.
1451fn 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
1459/// Reads one logical Mask batch, potentially spanning multiple loaded ranges.
1460///
1461/// Each [`MaskCursor`] chunk is safe to decode because it stays within loaded
1462/// pages. Gaps are crossed with [`ArrayReader::skip_records`], while decoded
1463/// arrays and their mask fragments remain buffered. Once `batch_size` selected
1464/// rows have accumulated, this consumes the underlying batch and filters it
1465/// once with the combined mask.
1466fn 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    /// Returns the next `RecordBatch` from the reader, or `None` if the reader
1550    /// has reached the end of the file.
1551    ///
1552    /// Returns `Result<Option<..>>` rather than `Option<Result<..>>` to
1553    /// simplify error handling with `?`
1554    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                    //Currently, when RowSelectors with row_count = 0 are included then its interpreted as end of reader.
1581                    //Fix is to skip such entries. See https://github.com/apache/arrow-rs/issues/2669
1582                    if front.row_count == 0 {
1583                        continue;
1584                    }
1585
1586                    // try to read record
1587                    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                            // if page row count less than batch_size we must set batch size to page row count.
1591                            // add check avoid dead loop
1592                            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    /// Returns the projected [`SchemaRef`] for reading the parquet file.
1619    ///
1620    /// Note that the schema metadata will be stripped here. See
1621    /// [`ParquetRecordBatchReaderBuilder::schema`] if the metadata is desired.
1622    fn schema(&self) -> SchemaRef {
1623        self.schema.clone()
1624    }
1625}
1626
1627impl ParquetRecordBatchReader {
1628    /// Create a new [`ParquetRecordBatchReader`] from the provided chunk reader
1629    ///
1630    /// See [`ParquetRecordBatchReaderBuilder`] for more options
1631    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    /// Create a new [`ParquetRecordBatchReader`] from the provided [`RowGroups`]
1638    ///
1639    /// Note: this is a low-level interface see [`ParquetRecordBatchReader::try_new`] for a
1640    /// higher-level interface for reading parquet data from a file
1641    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        // note metrics are not supported in this API
1648        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    /// Create a new [`ParquetRecordBatchReader`] that will read at most `batch_size` rows at
1666    /// a time from [`ArrayReader`] based on the configured `selection`. If `selection` is `None`
1667    /// all rows will be returned
1668    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        // Verify that the schema was correctly parsed
1770        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        // Verify that the metadata matches
1786        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        // test page encoding stats
1802        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        // test skipping all
1821        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        // test skipping all but one column and converting to mask
1838        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        // test skipping all
1863        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        // test skipping all but one column and converting to mask
1879        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        // Verify that the schema was correctly parsed
1907        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        // Verify that the schema was correctly parsed
1922        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(), // Cannot use &mut File (#1163)
1942            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        // Check they can be downcast to the correct type
2076        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        // Ensure can be downcast to the correct type
2141        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        // Ensure can be downcast to the correct type
2203        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        // Ensure can be downcast to the correct type
2238        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                // small-date64
2268                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                // big-date64
2276                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                // invalid-date64
2284                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        // Roundtrip should be successful when default writer used
2307        assert_eq!(default_ret, original);
2308
2309        // Only small-date64 should roundtrip successfully when coerce_types writer is used
2310        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        // Ensure both can be downcast to the correct type
2315        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            // Test without a specified ArrowType hint.
2383            (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            // Test other TimeUnits as ArrowType hints.
2389            (
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            // Test another timezone with TimeUnit as ArrowType hints.
2422            (
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        // int96_from_spark.parquet was written based on Spark's microsecond timestamps which trade
2452        // range for resolution compared to a nanosecond timestamp. We must provide a schema with
2453        // microsecond resolution for the Parquet reader to interpret these values correctly.
2454        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        // arrow-rs relies on the chrono library to convert between timestamps and strings, so
2487        // instead compare as Int64. The underlying type should be a PrimitiveArray of Int64
2488        // anyway, so this should be a zero-copy non-modifying cast.
2489
2490        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        // int96_from_spark.parquet was written based on Spark's microsecond timestamps which trade
2506        // range for resolution compared to a nanosecond timestamp. Without a provided schema, some
2507        // values when read as nanosecond resolution overflow and result in garbage values.
2508        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),  // Reads as nanosecond fine (note 3 extra 0s)
2525            Some(1704070800000000000),  // Reads as nanosecond fine (note 3 extra 0s)
2526            Some(-4852191831933722624), // Cannot be represented with nanos timestamp (year 9999)
2527            Some(1735599600000000000),  // Reads as nanosecond fine (note 3 extra 0s)
2528            None,
2529            Some(-4864435138808946688), // Cannot be represented with nanos timestamp (year 290000)
2530        ]));
2531
2532        // arrow-rs relies on the chrono library to convert between timestamps and strings, so
2533        // instead compare as Int64. The underlying type should be a PrimitiveArray of Int64
2534        // anyway, so this should be a zero-copy non-modifying cast.
2535
2536        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                // Cannot run full test suite as keys overflow, run small test instead
2618                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        // [[], [1], [2, 3], null, [4], null, [6, 7, 8]]
2747        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        // see https://github.com/apache/parquet-testing/pull/40
2816        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        // see https://github.com/apache/parquet-testing/pull/40
2848        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            // This file contains floats from a standard normal distribution
2884            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            // 0,1 are f16
2911            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            // 2,3 are f32
2920            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            // 4,5 are f64
2929            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            // 6,7 are i32
2938            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            // 8,9 are i64
2947            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            // 10,11 are FLBA(5)
2956            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            // 12,13 are FLBA(4) (decimal(7,3))
2965            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        // see https://github.com/apache/parquet-testing/pull/47
2980        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        // the `StructArrayReader` will check the definition and repetition levels of the first
3054        // child column in the struct to determine nullability for the struct. If the first
3055        // column's is being read by `ByteArrayDictionaryReader` we need to ensure that the
3056        // nullability is interpreted  correctly from the rep/def level buffers managed by the
3057        // buffers managed by this array reader.
3058
3059        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    /// Parameters for single_column_reader_test
3109    #[derive(Clone)]
3110    struct TestOptions {
3111        /// Number of row group to write to parquet (row group size =
3112        /// num_row_groups / num_rows)
3113        num_row_groups: usize,
3114        /// Total number of rows per row group
3115        num_rows: usize,
3116        /// Size of batches to read back
3117        record_batch_size: usize,
3118        /// Percentage of nulls in column or None if required
3119        null_percent: Option<usize>,
3120        /// Set write batch size
3121        ///
3122        /// This is the number of rows that are written at once to a page and
3123        /// therefore acts as a bound on the page granularity of a row group
3124        write_batch_size: usize,
3125        /// Maximum size of page in bytes
3126        max_data_page_size: usize,
3127        /// Maximum size of dictionary page in bytes
3128        max_dict_page_size: usize,
3129        /// Writer version
3130        writer_version: WriterVersion,
3131        /// Enabled statistics
3132        enabled_statistics: EnabledStatistics,
3133        /// Encoding
3134        encoding: Encoding,
3135        /// row selections and total selected row count
3136        row_selections: Option<(RowSelection, usize)>,
3137        /// row filter
3138        row_filter: Option<Vec<bool>>,
3139        /// limit
3140        limit: Option<usize>,
3141        /// offset
3142        offset: Option<usize>,
3143    }
3144
3145    /// Manually implement this to avoid printing entire contents of row_selections and row_filter
3146    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    /// Create a parquet file and then read it using
3290    /// `ParquetFileArrowReader` using a standard set of parameters
3291    /// `opts`.
3292    ///
3293    /// `rand_max` represents the maximum size of value to pass to to
3294    /// value generator
3295    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            // choose record_batch_batch (15) so batches cross row
3308            // group boundaries (50 rows in 2 row groups) cases.
3309            TestOptions::new(2, 100, 15),
3310            // choose record_batch_batch (5) so batches sometime fall
3311            // on row group boundaries and (25 rows in 3 row groups
3312            // --> row groups of 10, 10, and 5). Tests buffer
3313            // refilling edge cases.
3314            TestOptions::new(3, 25, 5),
3315            // Choose record_batch_size (25) so all batches fall
3316            // exactly on row group boundary (25). Tests buffer
3317            // refilling edge cases.
3318            TestOptions::new(4, 100, 25),
3319            // Set maximum page size so row groups have multiple pages
3320            TestOptions::new(3, 256, 73).with_max_data_page_size(128),
3321            // Set small dictionary page size to test dictionary fallback
3322            TestOptions::new(3, 256, 57).with_max_dict_page_size(128),
3323            // Test optional but with no nulls
3324            TestOptions::new(2, 256, 127).with_null_percent(0),
3325            // Test optional with nulls
3326            TestOptions::new(2, 256, 93).with_null_percent(25),
3327            // Test with limit of 0
3328            TestOptions::new(4, 100, 25).with_limit(0),
3329            // Test with limit of 50
3330            TestOptions::new(4, 100, 25).with_limit(50),
3331            // Test with limit equal to number of rows
3332            TestOptions::new(4, 100, 25).with_limit(10),
3333            // Test with limit larger than number of rows
3334            TestOptions::new(4, 100, 25).with_limit(101),
3335            // Test with limit + offset equal to number of rows
3336            TestOptions::new(4, 100, 25).with_offset(30).with_limit(20),
3337            // Test with limit + offset equal to number of rows
3338            TestOptions::new(4, 100, 25).with_offset(20).with_limit(80),
3339            // Test with limit + offset larger than number of rows
3340            TestOptions::new(4, 100, 25).with_offset(20).with_limit(81),
3341            // Test with no page-level statistics
3342            TestOptions::new(2, 256, 91)
3343                .with_null_percent(25)
3344                .with_enabled_statistics(EnabledStatistics::Chunk),
3345            // Test with no statistics
3346            TestOptions::new(2, 256, 91)
3347                .with_null_percent(25)
3348                .with_enabled_statistics(EnabledStatistics::None),
3349            // Test with all null
3350            TestOptions::new(2, 128, 91)
3351                .with_null_percent(100)
3352                .with_enabled_statistics(EnabledStatistics::None),
3353            // Test skip
3354
3355            // choose record_batch_batch (15) so batches cross row
3356            // group boundaries (50 rows in 2 row groups) cases.
3357            TestOptions::new(2, 100, 15).with_row_selections(),
3358            // choose record_batch_batch (5) so batches sometime fall
3359            // on row group boundaries and (25 rows in 3 row groups
3360            // --> row groups of 10, 10, and 5). Tests buffer
3361            // refilling edge cases.
3362            TestOptions::new(3, 25, 5).with_row_selections(),
3363            // Choose record_batch_size (25) so all batches fall
3364            // exactly on row group boundary (25). Tests buffer
3365            // refilling edge cases.
3366            TestOptions::new(4, 100, 25).with_row_selections(),
3367            // Set maximum page size so row groups have multiple pages
3368            TestOptions::new(3, 256, 73)
3369                .with_max_data_page_size(128)
3370                .with_row_selections(),
3371            // Set small dictionary page size to test dictionary fallback
3372            TestOptions::new(3, 256, 57)
3373                .with_max_dict_page_size(128)
3374                .with_row_selections(),
3375            // Test optional but with no nulls
3376            TestOptions::new(2, 256, 127)
3377                .with_null_percent(0)
3378                .with_row_selections(),
3379            // Test optional with nulls
3380            TestOptions::new(2, 256, 93)
3381                .with_null_percent(25)
3382                .with_row_selections(),
3383            // Test optional with nulls
3384            TestOptions::new(2, 256, 93)
3385                .with_null_percent(25)
3386                .with_row_selections()
3387                .with_limit(10),
3388            // Test optional with nulls
3389            TestOptions::new(2, 256, 93)
3390                .with_null_percent(25)
3391                .with_row_selections()
3392                .with_offset(20)
3393                .with_limit(10),
3394            // Test filter
3395
3396            // Test with row filter
3397            TestOptions::new(4, 100, 25).with_row_filter(),
3398            // Test with row selection and row filter
3399            TestOptions::new(4, 100, 25)
3400                .with_row_selections()
3401                .with_row_filter(),
3402            // Test with nulls and row filter
3403            TestOptions::new(2, 256, 93)
3404                .with_null_percent(25)
3405                .with_max_data_page_size(10)
3406                .with_row_filter(),
3407            // Test with nulls and row filter and small pages
3408            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            // Test with row selection and no offset index and small pages
3414            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    /// Create a parquet file and then read it using
3442    /// `ParquetFileArrowReader` using the parameters described in
3443    /// `opts`.
3444    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        // Print out options to facilitate debugging failures on CI
3456        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        //according to null_percent generate def_levels
3465        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        //generate random table data
3484        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(), // Cannot use &mut File (#1163)
3524            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                //get flatten table data
3559                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        // This particular test file has columns of struct types where there is
3706        // a column that has the same name as one of the struct fields
3707        // (see: ARROW-11452)
3708        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        // Tests for #1652 and #1654
3748        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    // same as test_read_structs but constructs projection mask via column names
3758    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 that we can handle the UNKNOWN logical type annotation on any physical type
3822    #[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            // write out a bunch of nulls
3847            column_writer
3848                .typed::<T>()
3849                .write_batch(&[], Some(&[0, 0, 0, 0]), None)
3850                .unwrap();
3851            column_writer.close().unwrap();
3852        }
3853
3854        // INT32
3855        write_nulls::<Int32Type>(&mut row_group_writer);
3856
3857        // INT64
3858        write_nulls::<Int64Type>(&mut row_group_writer);
3859
3860        // INT96
3861        write_nulls::<Int96Type>(&mut row_group_writer);
3862
3863        // BOOLEAN
3864        write_nulls::<BoolType>(&mut row_group_writer);
3865
3866        // FLOAT
3867        write_nulls::<FloatType>(&mut row_group_writer);
3868
3869        // DOUBLE
3870        write_nulls::<DoubleType>(&mut row_group_writer);
3871
3872        // BYTE_ARRAY
3873        write_nulls::<ByteArrayType>(&mut row_group_writer);
3874
3875        // FIXED_LEN_BYTE_ARRAY
3876        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            // Write using low-level parquet API (#1167)
3905            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(), // Cannot use &mut File (#1163)
3996            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        // First and second batch in same row group -> same dictionary
4026        assert_eq!(get_dict(&batches[0]), get_dict(&batches[1]));
4027        // Third batch spans row group -> computed dictionary
4028        assert_ne!(get_dict(&batches[1]), get_dict(&batches[2]));
4029        assert_ne!(get_dict(&batches[2]), get_dict(&batches[3]));
4030        // Fourth, fifth and sixth from same row group -> same dictionary
4031        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    /// Return parquet data with a single column of utf8 strings
4273    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        // write parquet file with non nullable strings
4278        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        // verify incompatible schemas error on read
4288        let parquet_data = utf8_parquet();
4289
4290        // Ask to read it back with an incompatible schema (int vs string)
4291        let input_schema: SchemaRef = Arc::new(Schema::new(vec![Field::new(
4292            "column1",
4293            arrow::datatypes::DataType::Int32,
4294            false,
4295        )]));
4296
4297        // read it back out
4298        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        // verify incompatible schemas error on read
4311        let parquet_data = utf8_parquet();
4312
4313        // Ask to read it back with an incompatible schema (nullability mismatch)
4314        let input_schema: SchemaRef = Arc::new(Schema::new(vec![Field::new(
4315            "column1",
4316            arrow::datatypes::DataType::Utf8,
4317            true,
4318        )]));
4319
4320        // read it back out
4321        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    /// Given a RecordBatch containing all the column data, return the expected batches given
4653    /// a `batch_size` and `selection`
4654    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        // Sanity check, all batches except the final should be the batch size
4701        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        // total row count 7300
4761        // 1. test selection len more than one page row count
4762        do_test(1000, 1000);
4763
4764        // 2. test selection len less than one page row count
4765        do_test(20, 20);
4766
4767        // 3. test selection_len less than batch_size
4768        do_test(20, 5);
4769
4770        // 4. test selection_len more than batch_size
4771        // If batch_size < selection_len
4772        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        // `alltypes_plain.parquet` only have 8 rows
4795        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            // `alltypes_tiny_pages.parquet` has page index
4815            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            // `alltypes_plain.parquet` doesn't have page index
4830            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            // Although `Vec<Vec<PageLoacation>>` of each row group is empty,
4838            // we should read the file successfully.
4839            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        // column 0
4865        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        // column 1
4872        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        // column 2
4879        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        // https://github.com/apache/parquet-testing/pull/18
4925        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    // This test is to ensure backward compatibility, it test 2 files containing the LZ4 CompressionCodec
4940    // but different algorithms: LZ4_HADOOP and LZ4_RAW.
4941    // 1. hadoop_lz4_compressed.parquet -> It is a file with LZ4 CompressionCodec which uses
4942    //    LZ4_HADOOP algorithm for compression.
4943    // 2. non_hadoop_lz4_compressed.parquet -> It is a file with LZ4 CompressionCodec which uses
4944    //    LZ4_RAW algorithm for compression. This fallback is done to keep backward compatibility with
4945    //    older parquet-cpp versions.
4946    //
4947    // For more information, check: https://github.com/apache/arrow-rs/issues/2988
4948    #[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        // First page contains 2 values but only 1 row
5080        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        // Precision <= 9 -> INT32
5130        // Precision <= 18 -> INT64
5131
5132        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        // Precision <= 9 -> INT32
5172        // Precision <= 18 -> INT64
5173        // Precision > 18 -> FIXED_LEN_BYTE_ARRAY
5174
5175        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        // message my_record {
5401        //     REQUIRED group a (LIST) {
5402        //         REPEATED group array (LIST) {
5403        //             REPEATED INT32 array;
5404        //         }
5405        //     }
5406        // }
5407        // should be read as list<list<int32>>
5408        let path = format!("{testdata}/old_list_structure.parquet");
5409        let test_file = File::open(path).unwrap();
5410
5411        // create expected ListArray
5412        let a_values = Int32Array::from(vec![1, 2, 3, 4]);
5413
5414        // Construct a buffer for value offsets, for the nested array: [[1, 2], [3, 4]]
5415        let a_value_offsets = arrow::buffer::Buffer::from([0, 2, 4].to_byte_slice());
5416
5417        // Construct a list array from the above two
5418        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        // grab first column
5436        let c0 = out.column(0);
5437        let c0arr = c0.as_any().downcast_ref::<ListArray>().unwrap();
5438        // get first row: [[1, 2], [3, 4]]
5439        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        // File schema:
5447        // message schema {
5448        //   required group my_map (MAP) {
5449        //     repeated group key_value {
5450        //       required int32 key;
5451        //       optional int32 value;
5452        //     }
5453        //   }
5454        //   required group my_map_no_v (MAP) {
5455        //     repeated group key_value {
5456        //       required int32 key;
5457        //     }
5458        //   }
5459        //   required group my_list (LIST) {
5460        //     repeated group list {
5461        //       required int32 element;
5462        //     }
5463        //   }
5464        // }
5465        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        // my_map_no_v and my_list columns should now be equivalent
5477        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        // Make a parquet file with 100 rows split across 2 row groups
5603        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        // write in 10 row batches as the size limits are enforced after each batch
5611        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        // read out with normal options
5626        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        // Now read, out of order row groups
5639        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        // Try to pass a regular field (not a virtual column) to with_virtual_columns
5694        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        // create a parquet file with 3 row groups, 2 rows each
5760        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        // write out 3 batches of 10 rows each
5893        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        // test row groups are read in reverse order
5910        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        // first row from rg 2
5922        assert_eq!(first_val, 20);
5923        // the last row from rg 0
5924        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        // assert_eq!(selection.row_count(), actual.num_rows());
5994        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}