Skip to main content

parquet/arrow/arrow_writer/
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 writer which writes arrow data into parquet data.
19
20use crate::column::chunker::ContentDefinedChunker;
21
22use bytes::Bytes;
23use std::io::Write;
24use std::slice::Iter;
25use std::sync::{Arc, Mutex};
26use std::vec::IntoIter;
27
28use arrow_array::cast::AsArray;
29use arrow_array::{ArrayRef, Int32Array, RecordBatch, RecordBatchWriter};
30use arrow_array::{PrimitiveArray, types::*};
31use arrow_schema::{
32    ArrowError, DataType as ArrowDataType, Field, IntervalUnit, SchemaRef, TimeUnit,
33};
34
35use super::schema::{add_encoded_arrow_schema_to_metadata, decimal_length_from_precision};
36
37use crate::arrow::ArrowSchemaConverter;
38use crate::arrow::arrow_writer::byte_array::ByteArrayEncoder;
39use crate::basic::PageType;
40use crate::column::page::{CompressedPage, PageWriteSpec, PageWriter};
41use crate::column::page_encryption::PageEncryptor;
42use crate::column::writer::encoder::ColumnValueEncoder;
43use crate::column::writer::{
44    ColumnCloseResult, ColumnWriter, GenericColumnWriter, get_column_writer,
45};
46use crate::data_type::{ByteArray, FixedLenByteArray};
47#[cfg(feature = "encryption")]
48use crate::encryption::encrypt::FileEncryptor;
49use crate::errors::{ParquetError, Result};
50use crate::file::metadata::{KeyValue, ParquetMetaData, RowGroupMetaData};
51use crate::file::properties::{WriterProperties, WriterPropertiesPtr};
52use crate::file::writer::{SerializedFileWriter, SerializedRowGroupWriter};
53use crate::parquet_thrift::{ThriftCompactOutputProtocol, WriteThrift};
54use crate::schema::types::{ColumnDescPtr, SchemaDescPtr, SchemaDescriptor};
55use levels::{ArrayLevels, calculate_array_levels};
56
57mod byte_array;
58mod levels;
59
60#[doc(inline)]
61pub use crate::column::page_store::{
62    InMemoryPageStore, InMemoryPageStoreFactory, PageKey, PageStore, PageStoreArgs,
63    PageStoreFactory,
64};
65
66/// Encodes [`RecordBatch`] to parquet
67///
68/// Writes Arrow `RecordBatch`es to a Parquet writer. Multiple [`RecordBatch`] will be encoded
69/// to the same row group, up to `max_row_group_size` rows. Any remaining rows will be
70/// flushed on close, leading the final row group in the output file to potentially
71/// contain fewer than `max_row_group_size` rows
72///
73/// # Example: Writing `RecordBatch`es
74/// ```
75/// # use std::sync::Arc;
76/// # use bytes::Bytes;
77/// # use arrow_array::{ArrayRef, Int64Array};
78/// # use arrow_array::RecordBatch;
79/// # use parquet::arrow::arrow_writer::ArrowWriter;
80/// # use parquet::arrow::arrow_reader::ParquetRecordBatchReader;
81/// let col = Arc::new(Int64Array::from_iter_values([1, 2, 3])) as ArrayRef;
82/// let to_write = RecordBatch::try_from_iter([("col", col)]).unwrap();
83///
84/// let mut buffer = Vec::new();
85/// let mut writer = ArrowWriter::try_new(&mut buffer, to_write.schema(), None).unwrap();
86/// writer.write(&to_write).unwrap();
87/// writer.close().unwrap();
88///
89/// let mut reader = ParquetRecordBatchReader::try_new(Bytes::from(buffer), 1024).unwrap();
90/// let read = reader.next().unwrap().unwrap();
91///
92/// assert_eq!(to_write, read);
93/// ```
94///
95/// # Memory Usage and Limiting
96///
97/// The nature of Parquet requires buffering of an entire row group before it can
98/// be flushed to the underlying writer. Data is mostly buffered in its encoded
99/// form, reducing memory usage. However, some data such as dictionary keys,
100/// large strings or very nested data may still result in non-trivial memory
101/// usage.
102///
103/// See Also:
104/// * [`ArrowWriter::memory_size`]: the current memory usage of the writer.
105/// * [`ArrowWriter::in_progress_size`]: Estimated size of the buffered row group,
106///
107/// Call [`Self::flush`] to trigger an early flush of a row group based on a
108/// memory threshold and/or global memory pressure. However,  smaller row groups
109/// result in higher metadata overheads, and thus may worsen compression ratios
110/// and query performance.
111///
112/// ```no_run
113/// # use std::io::Write;
114/// # use arrow_array::RecordBatch;
115/// # use parquet::arrow::ArrowWriter;
116/// # let mut writer: ArrowWriter<Vec<u8>> = todo!();
117/// # let batch: RecordBatch = todo!();
118/// writer.write(&batch).unwrap();
119/// // Trigger an early flush if anticipated size exceeds 1_000_000
120/// if writer.in_progress_size() > 1_000_000 {
121///     writer.flush().unwrap();
122/// }
123/// ```
124///
125/// ## Type Support
126///
127/// The writer supports writing all Arrow [`DataType`]s that have a direct mapping to
128/// Parquet types including  [`StructArray`] and [`ListArray`].
129///
130/// The following are not supported:
131///
132/// * [`IntervalMonthDayNanoArray`]: Parquet does not [support nanosecond intervals].
133///
134/// [`DataType`]: https://docs.rs/arrow/latest/arrow/datatypes/enum.DataType.html
135/// [`StructArray`]: https://docs.rs/arrow/latest/arrow/array/struct.StructArray.html
136/// [`ListArray`]: https://docs.rs/arrow/latest/arrow/array/type.ListArray.html
137/// [`IntervalMonthDayNanoArray`]: https://docs.rs/arrow/latest/arrow/array/type.IntervalMonthDayNanoArray.html
138/// [support nanosecond intervals]: https://github.com/apache/parquet-format/blob/master/LogicalTypes.md#interval
139///
140/// ## Type Compatibility
141/// The writer can write Arrow [`RecordBatch`]s that are logically equivalent. This means that for
142/// a  given column, the writer can accept multiple Arrow [`DataType`]s that contain the same
143/// value type.
144///
145/// For example, the following [`DataType`]s are all logically equivalent and can be written
146/// to the same column:
147/// * String, LargeString, StringView
148/// * Binary, LargeBinary, BinaryView
149///
150/// The writer can will also accept both native and dictionary encoded arrays if the dictionaries
151/// contain compatible values.
152/// ```
153/// # use std::sync::Arc;
154/// # use arrow_array::{DictionaryArray, LargeStringArray, RecordBatch, StringArray, UInt8Array};
155/// # use arrow_schema::{DataType, Field, Schema};
156/// # use parquet::arrow::arrow_writer::ArrowWriter;
157/// let record_batch1 = RecordBatch::try_new(
158///    Arc::new(Schema::new(vec![Field::new("col", DataType::LargeUtf8, false)])),
159///    vec![Arc::new(LargeStringArray::from_iter_values(vec!["a", "b"]))]
160///  )
161/// .unwrap();
162///
163/// let mut buffer = Vec::new();
164/// let mut writer = ArrowWriter::try_new(&mut buffer, record_batch1.schema(), None).unwrap();
165/// writer.write(&record_batch1).unwrap();
166///
167/// let record_batch2 = RecordBatch::try_new(
168///     Arc::new(Schema::new(vec![Field::new(
169///         "col",
170///         DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
171///          false,
172///     )])),
173///     vec![Arc::new(DictionaryArray::new(
174///          UInt8Array::from_iter_values(vec![0, 1]),
175///          Arc::new(StringArray::from_iter_values(vec!["b", "c"])),
176///      ))],
177///  )
178///  .unwrap();
179///  writer.write(&record_batch2).unwrap();
180///  writer.close();
181/// ```
182pub struct ArrowWriter<W: Write> {
183    /// Underlying Parquet writer
184    writer: SerializedFileWriter<W>,
185
186    /// The in-progress row group if any
187    in_progress: Option<ArrowRowGroupWriter>,
188
189    /// A copy of the Arrow schema.
190    ///
191    /// The schema is used to verify that each record batch written has the correct schema
192    arrow_schema: SchemaRef,
193
194    /// Creates new [`ArrowRowGroupWriter`] instances as required
195    row_group_writer_factory: ArrowRowGroupWriterFactory,
196
197    /// The maximum number of rows to write to each row group, or None for unlimited
198    max_row_group_row_count: Option<usize>,
199
200    /// The maximum size in bytes for a row group, or None for unlimited
201    max_row_group_bytes: Option<usize>,
202
203    /// CDC chunkers persisted across row groups (one per leaf column).
204    cdc_chunkers: Option<Vec<ContentDefinedChunker>>,
205}
206
207impl<W: Write + Send> std::fmt::Debug for ArrowWriter<W> {
208    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
209        let buffered_memory = self.in_progress_size();
210        f.debug_struct("ArrowWriter")
211            .field("writer", &self.writer)
212            .field("in_progress_size", &format_args!("{buffered_memory} bytes"))
213            .field("in_progress_rows", &self.in_progress_rows())
214            .field("arrow_schema", &self.arrow_schema)
215            .field("max_row_group_row_count", &self.max_row_group_row_count)
216            .field("max_row_group_bytes", &self.max_row_group_bytes)
217            .finish()
218    }
219}
220
221impl<W: Write + Send> ArrowWriter<W> {
222    /// Try to create a new Arrow writer
223    ///
224    /// The writer will fail if:
225    ///  * a `SerializedFileWriter` cannot be created from the ParquetWriter
226    ///  * the Arrow schema contains unsupported datatypes such as Unions
227    pub fn try_new(
228        writer: W,
229        arrow_schema: SchemaRef,
230        props: Option<WriterProperties>,
231    ) -> Result<Self> {
232        let options = ArrowWriterOptions::new().with_properties(props.unwrap_or_default());
233        Self::try_new_with_options(writer, arrow_schema, options)
234    }
235
236    /// Try to create a new Arrow writer with [`ArrowWriterOptions`].
237    ///
238    /// The writer will fail if:
239    ///  * a `SerializedFileWriter` cannot be created from the ParquetWriter
240    ///  * the Arrow schema contains unsupported datatypes such as Unions
241    pub fn try_new_with_options(
242        writer: W,
243        arrow_schema: SchemaRef,
244        options: ArrowWriterOptions,
245    ) -> Result<Self> {
246        let mut props = options.properties;
247
248        let schema = if let Some(parquet_schema) = options.schema_descr {
249            parquet_schema.clone()
250        } else {
251            let mut converter = ArrowSchemaConverter::new().with_coerce_types(props.coerce_types());
252            if let Some(schema_root) = &options.schema_root {
253                converter = converter.schema_root(schema_root);
254            }
255
256            converter.convert(&arrow_schema)?
257        };
258
259        if !options.skip_arrow_metadata {
260            // add serialized arrow schema
261            add_encoded_arrow_schema_to_metadata(&arrow_schema, &mut props);
262        }
263
264        let max_row_group_row_count = props.max_row_group_row_count();
265        let max_row_group_bytes = props.max_row_group_bytes();
266
267        let props_ptr = Arc::new(props);
268        let file_writer =
269            SerializedFileWriter::new(writer, schema.root_schema_ptr(), Arc::clone(&props_ptr))?;
270
271        let mut row_group_writer_factory =
272            ArrowRowGroupWriterFactory::new(&file_writer, arrow_schema.clone());
273        if let Some(page_store_factory) = options.page_store_factory {
274            row_group_writer_factory =
275                row_group_writer_factory.with_page_store_factory(page_store_factory);
276        }
277
278        let cdc_chunkers = props_ptr
279            .content_defined_chunking()
280            .map(|opts| {
281                file_writer
282                    .schema_descr()
283                    .columns()
284                    .iter()
285                    .map(|desc| ContentDefinedChunker::new(desc, opts))
286                    .collect::<Result<Vec<_>>>()
287            })
288            .transpose()?;
289
290        Ok(Self {
291            writer: file_writer,
292            in_progress: None,
293            arrow_schema,
294            row_group_writer_factory,
295            max_row_group_row_count,
296            max_row_group_bytes,
297            cdc_chunkers,
298        })
299    }
300
301    /// Returns metadata for any flushed row groups
302    pub fn flushed_row_groups(&self) -> &[RowGroupMetaData] {
303        self.writer.flushed_row_groups()
304    }
305
306    /// Estimated memory usage, in bytes, of this `ArrowWriter`
307    ///
308    /// This estimate is formed bu summing the values of
309    /// [`ArrowColumnWriter::memory_size`] all in progress columns.
310    pub fn memory_size(&self) -> usize {
311        match &self.in_progress {
312            Some(in_progress) => in_progress.writers.iter().map(|x| x.memory_size()).sum(),
313            None => 0,
314        }
315    }
316
317    /// Anticipated encoded size of the in progress row group.
318    ///
319    /// This estimate the row group size after being completely encoded is,
320    /// formed by summing the values of
321    /// [`ArrowColumnWriter::get_estimated_total_bytes`] for all in progress
322    /// columns.
323    pub fn in_progress_size(&self) -> usize {
324        match &self.in_progress {
325            Some(in_progress) => in_progress
326                .writers
327                .iter()
328                .map(|x| x.get_estimated_total_bytes())
329                .sum(),
330            None => 0,
331        }
332    }
333
334    /// Returns the number of rows buffered in the in progress row group
335    pub fn in_progress_rows(&self) -> usize {
336        self.in_progress
337            .as_ref()
338            .map(|x| x.buffered_rows)
339            .unwrap_or_default()
340    }
341
342    /// Returns the number of bytes written by this instance
343    pub fn bytes_written(&self) -> usize {
344        self.writer.bytes_written()
345    }
346
347    /// Encodes the provided [`RecordBatch`]
348    ///
349    /// If this would cause the current row group to exceed [`WriterProperties::max_row_group_row_count`]
350    /// rows or [`WriterProperties::max_row_group_bytes`] bytes, the contents of `batch` will be
351    /// written to one or more row groups such that limits are respected.
352    ///
353    /// If both limits are `None`, all data is written to a single row group.
354    /// If one limit is set, that limit is respected.
355    /// If both limits are set, the lower bound (whichever triggers first) is respected.
356    ///
357    /// This will fail if the `batch`'s schema does not match the writer's schema.
358    pub fn write(&mut self, batch: &RecordBatch) -> Result<()> {
359        if batch.num_rows() == 0 {
360            return Ok(());
361        }
362
363        let in_progress = match &mut self.in_progress {
364            Some(in_progress) => in_progress,
365            x => x.insert(
366                self.row_group_writer_factory
367                    .create_row_group_writer(self.writer.flushed_row_groups().len())?,
368            ),
369        };
370
371        if let Some(max_rows) = self.max_row_group_row_count
372            && in_progress.buffered_rows + batch.num_rows() > max_rows
373        {
374            let to_write = max_rows - in_progress.buffered_rows;
375            let a = batch.slice(0, to_write);
376            let b = batch.slice(to_write, batch.num_rows() - to_write);
377            self.write(&a)?;
378            return self.write(&b);
379        }
380
381        // Check byte limit: if we have buffered data, use measured average row size
382        // to split batch proactively before exceeding byte limit
383        if let Some(max_bytes) = self.max_row_group_bytes
384            && in_progress.buffered_rows > 0
385        {
386            let current_bytes = in_progress.get_estimated_total_bytes();
387
388            if current_bytes >= max_bytes {
389                self.flush()?;
390                return self.write(batch);
391            }
392
393            if let Some(avg_row_bytes) = current_bytes
394                .checked_div(in_progress.buffered_rows)
395                .filter(|avg_row_bytes| *avg_row_bytes > 0)
396            {
397                // At this point, `current_bytes < max_bytes` (checked above)
398                let remaining_bytes = max_bytes - current_bytes;
399                let rows_that_fit = remaining_bytes.checked_div(avg_row_bytes).unwrap_or(0);
400
401                if batch.num_rows() > rows_that_fit {
402                    if rows_that_fit > 0 {
403                        let a = batch.slice(0, rows_that_fit);
404                        let b = batch.slice(rows_that_fit, batch.num_rows() - rows_that_fit);
405                        self.write(&a)?;
406                        return self.write(&b);
407                    } else {
408                        self.flush()?;
409                        return self.write(batch);
410                    }
411                }
412            }
413        }
414
415        match self.cdc_chunkers.as_mut() {
416            Some(chunkers) => in_progress.write_with_chunkers(batch, chunkers)?,
417            None => in_progress.write(batch)?,
418        }
419
420        let should_flush = self
421            .max_row_group_row_count
422            .is_some_and(|max| in_progress.buffered_rows >= max)
423            || self
424                .max_row_group_bytes
425                .is_some_and(|max| in_progress.get_estimated_total_bytes() >= max);
426
427        if should_flush {
428            self.flush()?
429        }
430        Ok(())
431    }
432
433    /// Writes the given buf bytes to the internal buffer.
434    ///
435    /// It's safe to use this method to write data to the underlying writer,
436    /// because it will ensure that the buffering and byte‐counting layers are used.
437    pub fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
438        self.writer.write_all(buf)
439    }
440
441    /// Flushes underlying writer
442    pub fn sync(&mut self) -> std::io::Result<()> {
443        self.writer.flush()
444    }
445
446    /// Flushes all buffered rows into a new row group
447    ///
448    /// Note the underlying writer is not flushed with this call.
449    /// If this is a desired behavior, please call [`ArrowWriter::sync`].
450    pub fn flush(&mut self) -> Result<()> {
451        let in_progress = match self.in_progress.take() {
452            Some(in_progress) => in_progress,
453            None => return Ok(()),
454        };
455
456        let mut row_group_writer = self.writer.next_row_group()?;
457        for chunk in in_progress.close()? {
458            chunk.append_to_row_group(&mut row_group_writer)?;
459        }
460        row_group_writer.close()?;
461        Ok(())
462    }
463
464    /// Additional [`KeyValue`] metadata to be written in addition to those from [`WriterProperties`]
465    ///
466    /// This method provide a way to append kv_metadata after write RecordBatch
467    pub fn append_key_value_metadata(&mut self, kv_metadata: KeyValue) {
468        self.writer.append_key_value_metadata(kv_metadata)
469    }
470
471    /// Returns a reference to the underlying writer.
472    pub fn inner(&self) -> &W {
473        self.writer.inner()
474    }
475
476    /// Returns a mutable reference to the underlying writer.
477    ///
478    /// **Warning**: if you write directly to this writer, you will skip
479    /// the `TrackedWrite` buffering and byte‐counting layers. That’ll cause
480    /// the file footer’s recorded offsets and sizes to diverge from reality,
481    /// resulting in an unreadable or corrupted Parquet file.
482    ///
483    /// If you want to write safely to the underlying writer, use [`Self::write_all`].
484    pub fn inner_mut(&mut self) -> &mut W {
485        self.writer.inner_mut()
486    }
487
488    /// Flushes any outstanding data and returns the underlying writer.
489    pub fn into_inner(mut self) -> Result<W> {
490        self.flush()?;
491        self.writer.into_inner()
492    }
493
494    /// Close and finalize the underlying Parquet writer
495    ///
496    /// Unlike [`Self::close`] this does not consume self
497    ///
498    /// Attempting to write after calling finish will result in an error
499    pub fn finish(&mut self) -> Result<ParquetMetaData> {
500        self.flush()?;
501        self.writer.finish()
502    }
503
504    /// Close and finalize the underlying Parquet writer
505    pub fn close(mut self) -> Result<ParquetMetaData> {
506        self.finish()
507    }
508
509    /// Converts this writer into a lower-level [`SerializedFileWriter`] and [`ArrowRowGroupWriterFactory`].
510    ///
511    /// Flushes any outstanding data before returning.
512    ///
513    /// This can be useful to provide more control over how files are written, for example
514    /// to write columns in parallel. See the example on [`ArrowColumnWriter`].
515    pub fn into_serialized_writer(
516        mut self,
517    ) -> Result<(SerializedFileWriter<W>, ArrowRowGroupWriterFactory)> {
518        self.flush()?;
519        Ok((self.writer, self.row_group_writer_factory))
520    }
521}
522
523impl<W: Write + Send> RecordBatchWriter for ArrowWriter<W> {
524    fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
525        self.write(batch).map_err(|e| e.into())
526    }
527
528    fn close(self) -> std::result::Result<(), ArrowError> {
529        self.close()?;
530        Ok(())
531    }
532}
533
534/// Arrow-specific configuration settings for writing parquet files.
535///
536/// See [`ArrowWriter`] for how to configure the writer.
537#[derive(Debug, Clone, Default)]
538pub struct ArrowWriterOptions {
539    properties: WriterProperties,
540    skip_arrow_metadata: bool,
541    schema_root: Option<String>,
542    schema_descr: Option<SchemaDescriptor>,
543    page_store_factory: Option<Arc<dyn PageStoreFactory>>,
544}
545
546impl ArrowWriterOptions {
547    /// Creates a new [`ArrowWriterOptions`] with the default settings.
548    pub fn new() -> Self {
549        Self::default()
550    }
551
552    /// Sets the [`WriterProperties`] for writing parquet files.
553    pub fn with_properties(self, properties: WriterProperties) -> Self {
554        Self { properties, ..self }
555    }
556
557    /// Sets the [`PageStoreFactory`] used to buffer completed pages while a row
558    /// group is being written.
559    ///
560    /// The default implementation ([`InMemoryPageStore`]) buffers all completed
561    /// pages on the heap until the row group is flushed, so peak write memory
562    /// grows with the row group size. Using this API, pages can be spilled to a
563    /// file or object storage instead, reducing peak write memory substantially
564    /// at the expense of an extra write to and read from secondary storage.
565    ///
566    /// # Example: spilling pages to a temp file
567    ///
568    /// A simple spilling backend uses one temp file per column chunk; `put`
569    /// appends the page and `take` reads it back.
570    ///
571    /// ```
572    /// # use std::fs::File;
573    /// # use std::io::{Read, Seek, SeekFrom, Write};
574    /// # use std::sync::Arc;
575    /// # use bytes::Bytes;
576    /// # use arrow_array::{ArrayRef, Int64Array, RecordBatch};
577    /// # use parquet::arrow::arrow_writer::{
578    /// #     ArrowWriter, ArrowWriterOptions, PageKey, PageStore, PageStoreArgs, PageStoreFactory,
579    /// # };
580    /// # use parquet::arrow::arrow_reader::ParquetRecordBatchReader;
581    /// # use parquet::errors::Result;
582    /// struct TempFilePageStore {
583    ///     file: File,
584    ///     /// Total size of the file
585    ///     end: u64,
586    ///     /// Location of pages: (offset, len)
587    ///     locs: Vec<(u64, usize)>,
588    /// }
589    ///
590    /// impl PageStore for TempFilePageStore {
591    ///     fn put(&mut self, value: Bytes) -> Result<PageKey> {
592    ///         // Append to the end of the file
593    ///         self.file.seek(SeekFrom::Start(self.end))?;
594    ///         self.file.write_all(&value)?;
595    ///         let key = PageKey::new(self.locs.len() as u64);
596    ///         self.locs.push((self.end, value.len()));
597    ///         self.end += value.len() as u64;
598    ///         Ok(key)
599    ///     }
600    ///
601    ///     fn take(&mut self, key: PageKey) -> Result<Bytes> {
602    ///         let (offset, len) = self.locs[key.get() as usize];
603    ///         let mut buf = vec![0u8; len];
604    ///         self.file.seek(SeekFrom::Start(offset))?;
605    ///         self.file.read_exact(&mut buf)?;
606    ///         Ok(Bytes::from(buf))
607    ///     }
608    /// }
609    ///
610    /// /// Factory for creating [`TempFilePageStore`]
611    /// #[derive(Debug)]
612    /// struct TempFilePageStoreFactory;
613    ///
614    /// impl PageStoreFactory for TempFilePageStoreFactory {
615    ///     fn create(&self, args: &PageStoreArgs<'_>) -> Result<Box<dyn PageStore>> {
616    ///         // `args` exposes the column index and descriptor (physical/logical
617    ///         // type, path), so a real backend might choose to spill only large columns.
618    ///         let _ = (args.column_index(), args.column_descriptor());
619    ///         Ok(Box::new(TempFilePageStore {
620    ///             file: tempfile::tempfile()?, // temp file is cleaned on drop
621    ///             end: 0,
622    ///             locs: Vec::new(),
623    ///         }))
624    ///     }
625    /// }
626    /// // write 1000 integers
627    /// let col = Arc::new(Int64Array::from_iter_values(0..1000)) as ArrayRef;
628    /// let to_write = RecordBatch::try_from_iter([("col", col)]).unwrap();
629    ///
630    /// let options =
631    ///     ArrowWriterOptions::new().with_page_store_factory(Arc::new(TempFilePageStoreFactory));
632    /// let mut buffer = Vec::new();
633    /// let mut writer =
634    ///     ArrowWriter::try_new_with_options(&mut buffer, to_write.schema(), options).unwrap();
635    /// writer.write(&to_write).unwrap();
636    /// writer.close().unwrap();
637    ///
638    /// // buffer now holds valid Parquet data, which can be read as normal:
639    /// let mut reader = ParquetRecordBatchReader::try_new(Bytes::from(buffer), 1024).unwrap();
640    /// assert_eq!(to_write, reader.next().unwrap().unwrap());
641    /// ```
642    pub fn with_page_store_factory(self, page_store_factory: Arc<dyn PageStoreFactory>) -> Self {
643        Self {
644            page_store_factory: Some(page_store_factory),
645            ..self
646        }
647    }
648
649    /// Skip encoding the embedded arrow metadata (defaults to `false`)
650    ///
651    /// Parquet files generated by the [`ArrowWriter`] contain embedded arrow schema
652    /// by default.
653    ///
654    /// Set `skip_arrow_metadata` to true, to skip encoding the embedded metadata.
655    pub fn with_skip_arrow_metadata(self, skip_arrow_metadata: bool) -> Self {
656        Self {
657            skip_arrow_metadata,
658            ..self
659        }
660    }
661
662    /// Set the name of the root parquet schema element (defaults to `"arrow_schema"`)
663    pub fn with_schema_root(self, schema_root: String) -> Self {
664        Self {
665            schema_root: Some(schema_root),
666            ..self
667        }
668    }
669
670    /// Explicitly specify the Parquet schema to be used
671    ///
672    /// If omitted (the default), the [`ArrowSchemaConverter`] is used to compute the
673    /// Parquet [`SchemaDescriptor`]. This may be used When the [`SchemaDescriptor`] is
674    /// already known or must be calculated using custom logic.
675    pub fn with_parquet_schema(self, schema_descr: SchemaDescriptor) -> Self {
676        Self {
677            schema_descr: Some(schema_descr),
678            ..self
679        }
680    }
681}
682
683/// A single column chunk produced by [`ArrowColumnWriter`].
684///
685/// Holds the serialized page blobs (each page's header ‖ compressed data, in
686/// write order) in a [`PageStore`], plus the handles needed to read them back,
687/// in order, when the chunk is spliced into the output file.
688struct ArrowColumnChunkData {
689    length: usize,
690    store: Box<dyn PageStore>,
691    keys: Vec<PageKey>,
692    /// Handles to the dictionary page's blobs (header then data) in the store.
693    ///
694    /// A dictionary page is produced at most once and bounded by
695    /// `dict_page_size_limit`, but it must be written *first* in the chunk even
696    /// though the data pages reach the writer before it (see
697    /// [`PageWriter::defers_dictionary_ordering`]). Its header and data are `put`
698    /// into the store like any other page — which keeps the store uniform, and
699    /// lets an oversized dictionary page spill — and their handles are held apart
700    /// so they can be emitted ahead of the data pages at splice.
701    /// Empty for non-dictionary columns.
702    dictionary_keys: Vec<PageKey>,
703    /// Serialized length of the dictionary page (0 if there is none), recorded
704    /// so the data pages can be shifted past it when offsets are rewritten to a
705    /// dictionary-first layout at splice.
706    dictionary_len: usize,
707}
708
709impl ArrowColumnChunkData {
710    fn new(store: Box<dyn PageStore>) -> Self {
711        Self {
712            length: 0,
713            store,
714            keys: Vec::new(),
715            dictionary_keys: Vec::new(),
716            dictionary_len: 0,
717        }
718    }
719
720    /// Append a data-page blob to the store, recording its handle in write
721    /// order.
722    fn push(&mut self, value: Bytes) -> Result<()> {
723        let key = self.store.put(value)?;
724        self.keys.push(key);
725        Ok(())
726    }
727
728    /// Store a dictionary-page blob (header or data) in the page store,
729    /// recording its handle (emitted first at splice) and accumulating its
730    /// serialized length.
731    fn push_dictionary(&mut self, value: Bytes) -> Result<()> {
732        self.dictionary_len += value.len();
733        let key = self.store.put(value)?;
734        self.dictionary_keys.push(key);
735        Ok(())
736    }
737
738    /// Bytes this chunk currently holds on the heap: whatever the store keeps
739    /// resident (zero for a spilling backend).
740    fn memory_size(&self) -> usize {
741        self.store.memory_size()
742    }
743}
744
745/// A streaming iterator over one column chunk's buffered page blobs, in final
746/// file order: the dictionary page (if any) first, then the data pages.
747///
748/// Each blob is taken back out of the [`PageStore`] *as it is
749/// consumed* and released immediately afterwards, so splicing a chunk into the
750/// output file never materializes more than a single page in memory at a time.
751/// This is what keeps the splice phase within the memory bound for a spilling
752/// backend (an in-memory store already holds the bytes, so it is unaffected).
753struct StreamingColumnChunkPages {
754    store: Box<dyn PageStore>,
755    /// Page handles in final file order: the dictionary page first (if any),
756    /// then the data pages.
757    keys: IntoIter<PageKey>,
758}
759
760impl StreamingColumnChunkPages {
761    fn new(data: ArrowColumnChunkData) -> Self {
762        // The dictionary page must be emitted first, ahead of the data pages,
763        // even though it was the last page produced.
764        let keys = if data.dictionary_keys.is_empty() {
765            data.keys
766        } else {
767            let mut keys = Vec::with_capacity(data.dictionary_keys.len() + data.keys.len());
768            keys.extend(data.dictionary_keys);
769            keys.extend(data.keys);
770            keys
771        };
772        Self {
773            store: data.store,
774            keys: keys.into_iter(),
775        }
776    }
777}
778
779impl Iterator for StreamingColumnChunkPages {
780    type Item = Result<Bytes>;
781
782    fn next(&mut self) -> Option<Self::Item> {
783        let key = self.keys.next()?;
784        Some(self.store.take(key))
785    }
786}
787
788/// A shared [`ArrowColumnChunkData`]
789///
790/// This allows it to be owned by [`ArrowPageWriter`] whilst allowing access via
791/// [`ArrowRowGroupWriter`] on flush, without requiring self-referential borrows
792type SharedColumnChunk = Arc<Mutex<ArrowColumnChunkData>>;
793
794struct ArrowPageWriter {
795    buffer: SharedColumnChunk,
796    #[cfg(feature = "encryption")]
797    page_encryptor: Option<PageEncryptor>,
798}
799
800impl ArrowPageWriter {
801    /// Create a page writer that buffers completed pages in `store`.
802    fn new(store: Box<dyn PageStore>) -> Self {
803        Self {
804            buffer: Arc::new(Mutex::new(ArrowColumnChunkData::new(store))),
805            #[cfg(feature = "encryption")]
806            page_encryptor: None,
807        }
808    }
809
810    #[cfg(feature = "encryption")]
811    pub fn with_encryptor(mut self, page_encryptor: Option<PageEncryptor>) -> Self {
812        self.page_encryptor = page_encryptor;
813        self
814    }
815
816    #[cfg(feature = "encryption")]
817    fn page_encryptor_mut(&mut self) -> Option<&mut PageEncryptor> {
818        self.page_encryptor.as_mut()
819    }
820
821    #[cfg(not(feature = "encryption"))]
822    fn page_encryptor_mut(&mut self) -> Option<&mut PageEncryptor> {
823        None
824    }
825}
826
827impl PageWriter for ArrowPageWriter {
828    fn write_page(&mut self, page: CompressedPage) -> Result<PageWriteSpec> {
829        let page = match self.page_encryptor_mut() {
830            Some(page_encryptor) => page_encryptor.encrypt_compressed_page(page)?,
831            None => page,
832        };
833
834        let page_header = page.to_thrift_header()?;
835        let header = {
836            let mut header = Vec::with_capacity(1024);
837
838            match self.page_encryptor_mut() {
839                Some(page_encryptor) => {
840                    page_encryptor.encrypt_page_header(&page_header, &mut header)?;
841                    if page.compressed_page().is_data_page() {
842                        page_encryptor.increment_page();
843                    }
844                }
845                None => {
846                    let mut protocol = ThriftCompactOutputProtocol::new(&mut header);
847                    page_header.write_thrift(&mut protocol)?;
848                }
849            };
850
851            Bytes::from(header)
852        };
853
854        let mut buf = self.buffer.try_lock().unwrap();
855
856        let data = page.compressed_page().buffer().clone();
857        let compressed_size = data.len() + header.len();
858
859        let mut spec = PageWriteSpec::new();
860        spec.page_type = page.page_type();
861        spec.num_values = page.num_values();
862        spec.uncompressed_size = page.uncompressed_size() + header.len();
863        spec.offset = buf.length as u64;
864        spec.compressed_size = compressed_size;
865        spec.bytes_written = compressed_size as u64;
866
867        buf.length += compressed_size;
868        if spec.page_type == PageType::DICTIONARY_PAGE {
869            // Recorded apart from the data pages so it is emitted first at
870            // splice — see `ArrowColumnChunkData::dictionary_keys`.
871            buf.push_dictionary(header)?;
872            buf.push_dictionary(data)?;
873        } else {
874            buf.push(header)?;
875            buf.push(data)?;
876        }
877
878        Ok(spec)
879    }
880
881    fn defers_dictionary_ordering(&self) -> bool {
882        // The Arrow chunk is buffered in full and spliced at row-group flush, so
883        // data pages may be accepted before the dictionary page and reordered
884        // then. This lets `GenericColumnWriter` stream dictionary-column data
885        // pages straight through instead of buffering them in memory.
886        true
887    }
888
889    fn buffered_memory_size(&self) -> usize {
890        // Only what is actually resident: a spilling store reports ~0 here even
891        // though the chunk's bytes have all passed through it.
892        self.buffer.try_lock().unwrap().memory_size()
893    }
894
895    fn close(&mut self) -> Result<()> {
896        Ok(())
897    }
898}
899
900/// A leaf column that can be encoded by [`ArrowColumnWriter`]
901#[derive(Debug)]
902pub struct ArrowLeafColumn(ArrayLevels);
903
904/// Computes the [`ArrowLeafColumn`] for a potentially nested [`ArrayRef`]
905///
906/// This function can be used to encode individual columns in parallel.
907/// See example on [`ArrowColumnWriter`]
908pub fn compute_leaves(field: &Field, array: &ArrayRef) -> Result<Vec<ArrowLeafColumn>> {
909    let levels = calculate_array_levels(array, field)?;
910    Ok(levels.into_iter().map(ArrowLeafColumn).collect())
911}
912
913/// The data for a single column chunk, see [`ArrowColumnWriter`]
914pub struct ArrowColumnChunk {
915    data: ArrowColumnChunkData,
916    close: ColumnCloseResult,
917}
918
919impl std::fmt::Debug for ArrowColumnChunk {
920    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
921        f.debug_struct("ArrowColumnChunk")
922            .field("length", &self.data.length)
923            .finish_non_exhaustive()
924    }
925}
926
927impl ArrowColumnChunk {
928    /// Returns the [`ColumnCloseResult`] produced when the chunk was closed.
929    ///
930    /// Exposes encoding information, collected statistics, and the optional
931    /// [`ColumnIndexMetaData`](crate::file::page_index::column_index::ColumnIndexMetaData)
932    /// / [`OffsetIndexMetaData`](crate::file::page_index::offset_index::OffsetIndexMetaData)
933    /// gathered for the column chunk.
934    pub fn close(&self) -> &ColumnCloseResult {
935        &self.close
936    }
937
938    /// Returns a mutable reference to the [`ColumnCloseResult`].
939    ///
940    /// This allows callers to mutate the close result before the chunk is
941    /// appended to a row group — for example, clearing `column_index` or
942    /// `bloom_filter` based on a dynamic rule that inspects the encodings and
943    /// collected page statistics.
944    pub fn close_mut(&mut self) -> &mut ColumnCloseResult {
945        &mut self.close
946    }
947
948    /// Splices this column's buffered pages into the row group, streaming them
949    /// back out of the [`PageStore`] one page at a time.
950    pub fn append_to_row_group<W: Write + Send>(
951        self,
952        writer: &mut SerializedRowGroupWriter<'_, W>,
953    ) -> Result<()> {
954        let ArrowColumnChunk { data, close } = self;
955
956        // The dictionary page is produced *after* the data pages on this path (so
957        // they can stream straight through) but must be written *first*, so move
958        // it ahead of the data pages in the recorded offsets before the splice.
959        let close = close.update_dictionary_location(data.dictionary_len)?;
960
961        let pages = StreamingColumnChunkPages::new(data);
962        writer.append_column_from_pages(pages, close)
963    }
964}
965
966/// Encodes [`ArrowLeafColumn`] to [`ArrowColumnChunk`]
967///
968/// `ArrowColumnWriter` instances can be created using an [`ArrowRowGroupWriterFactory`];
969///
970/// Note: This is a low-level interface for applications that require
971/// fine-grained control of encoding (e.g. encoding using multiple threads),
972/// see [`ArrowWriter`] for a higher-level interface
973///
974/// # Example: Encoding two Arrow Array's in Parallel
975/// ```
976/// // The arrow schema
977/// # use std::sync::Arc;
978/// # use arrow_array::*;
979/// # use arrow_schema::*;
980/// # use parquet::arrow::ArrowSchemaConverter;
981/// # use parquet::arrow::arrow_writer::{compute_leaves, ArrowColumnChunk, ArrowLeafColumn, ArrowRowGroupWriterFactory};
982/// # use parquet::file::properties::WriterProperties;
983/// # use parquet::file::writer::{SerializedFileWriter, SerializedRowGroupWriter};
984/// #
985/// let schema = Arc::new(Schema::new(vec![
986///     Field::new("i32", DataType::Int32, false),
987///     Field::new("f32", DataType::Float32, false),
988/// ]));
989///
990/// // Compute the parquet schema
991/// let props = Arc::new(WriterProperties::default());
992/// let parquet_schema = ArrowSchemaConverter::new()
993///   .with_coerce_types(props.coerce_types())
994///   .convert(&schema)
995///   .unwrap();
996///
997/// // Create parquet writer
998/// let root_schema = parquet_schema.root_schema_ptr();
999/// // write to memory in the example, but this could be a File
1000/// let mut out = Vec::with_capacity(1024);
1001/// let mut writer = SerializedFileWriter::new(&mut out, root_schema, props.clone())
1002///   .unwrap();
1003///
1004/// // Create a factory for building Arrow column writers
1005/// let row_group_factory = ArrowRowGroupWriterFactory::new(&writer, Arc::clone(&schema));
1006/// // Create column writers for the 0th row group
1007/// let col_writers = row_group_factory.create_column_writers(0).unwrap();
1008///
1009/// // Spawn a worker thread for each column
1010/// //
1011/// // Note: This is for demonstration purposes, a thread-pool e.g. rayon or tokio, would be better.
1012/// // The `map` produces an iterator of type `tuple of (thread handle, send channel)`.
1013/// let mut workers: Vec<_> = col_writers
1014///     .into_iter()
1015///     .map(|mut col_writer| {
1016///         let (send, recv) = std::sync::mpsc::channel::<ArrowLeafColumn>();
1017///         let handle = std::thread::spawn(move || {
1018///             // receive Arrays to encode via the channel
1019///             for col in recv {
1020///                 col_writer.write(&col)?;
1021///             }
1022///             // once the input is complete, close the writer
1023///             // to return the newly created ArrowColumnChunk
1024///             col_writer.close()
1025///         });
1026///         (handle, send)
1027///     })
1028///     .collect();
1029///
1030/// // Start row group
1031/// let mut row_group_writer: SerializedRowGroupWriter<'_, _> = writer
1032///   .next_row_group()
1033///   .unwrap();
1034///
1035/// // Create some example input columns to encode
1036/// let to_write = vec![
1037///     Arc::new(Int32Array::from_iter_values([1, 2, 3])) as _,
1038///     Arc::new(Float32Array::from_iter_values([1., 45., -1.])) as _,
1039/// ];
1040///
1041/// // Send the input columns to the workers
1042/// let mut worker_iter = workers.iter_mut();
1043/// for (arr, field) in to_write.iter().zip(&schema.fields) {
1044///     for leaves in compute_leaves(field, arr).unwrap() {
1045///         worker_iter.next().unwrap().1.send(leaves).unwrap();
1046///     }
1047/// }
1048///
1049/// // Wait for the workers to complete encoding, and append
1050/// // the resulting column chunks to the row group (and the file)
1051/// for (handle, send) in workers {
1052///     drop(send); // Drop send side to signal termination
1053///     // wait for the worker to send the completed chunk
1054///     let chunk: ArrowColumnChunk = handle.join().unwrap().unwrap();
1055///     chunk.append_to_row_group(&mut row_group_writer).unwrap();
1056/// }
1057/// // Close the row group which writes to the underlying file
1058/// row_group_writer.close().unwrap();
1059///
1060/// let metadata = writer.close().unwrap();
1061/// assert_eq!(metadata.file_metadata().num_rows(), 3);
1062/// ```
1063pub struct ArrowColumnWriter {
1064    writer: ArrowColumnWriterImpl,
1065    chunk: SharedColumnChunk,
1066}
1067
1068impl std::fmt::Debug for ArrowColumnWriter {
1069    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1070        f.debug_struct("ArrowColumnWriter").finish_non_exhaustive()
1071    }
1072}
1073
1074enum ArrowColumnWriterImpl {
1075    ByteArray(GenericColumnWriter<'static, ByteArrayEncoder>),
1076    Column(ColumnWriter<'static>),
1077}
1078
1079impl ArrowColumnWriter {
1080    /// Write an [`ArrowLeafColumn`]
1081    pub fn write(&mut self, col: &ArrowLeafColumn) -> Result<()> {
1082        self.write_internal(&col.0)
1083    }
1084
1085    /// Write with content-defined chunking, inserting page flushes at chunk boundaries.
1086    fn write_with_chunker(
1087        &mut self,
1088        col: &ArrowLeafColumn,
1089        chunker: &mut ContentDefinedChunker,
1090    ) -> Result<()> {
1091        let levels = &col.0;
1092        let chunks = chunker.get_arrow_chunks(
1093            levels.def_level_data().as_ref(),
1094            levels.rep_level_data().as_ref(),
1095            levels.array(),
1096        )?;
1097
1098        let num_chunks = chunks.len();
1099        for (i, chunk) in chunks.iter().enumerate() {
1100            let chunk_levels = levels.slice_for_chunk(chunk);
1101            self.write_internal(&chunk_levels)?;
1102
1103            // Add a page break after each chunk except the last
1104            if i + 1 < num_chunks {
1105                match &mut self.writer {
1106                    ArrowColumnWriterImpl::Column(c) => c.add_data_page()?,
1107                    ArrowColumnWriterImpl::ByteArray(c) => c.add_data_page()?,
1108                }
1109            }
1110        }
1111        Ok(())
1112    }
1113
1114    fn write_internal(&mut self, levels: &ArrayLevels) -> Result<()> {
1115        match &mut self.writer {
1116            ArrowColumnWriterImpl::Column(c) => {
1117                let leaf = levels.array();
1118                match leaf.as_any_dictionary_opt() {
1119                    Some(dictionary) => {
1120                        let materialized =
1121                            arrow_select::take::take(dictionary.values(), dictionary.keys(), None)?;
1122                        write_leaf(c, &materialized, levels)?
1123                    }
1124                    None => write_leaf(c, leaf, levels)?,
1125                };
1126            }
1127            ArrowColumnWriterImpl::ByteArray(c) => {
1128                write_primitive(c, levels.array().as_ref(), levels)?;
1129            }
1130        }
1131        Ok(())
1132    }
1133
1134    /// Close this column returning the written [`ArrowColumnChunk`]
1135    pub fn close(self) -> Result<ArrowColumnChunk> {
1136        let close = match self.writer {
1137            ArrowColumnWriterImpl::ByteArray(c) => c.close()?,
1138            ArrowColumnWriterImpl::Column(c) => c.close()?,
1139        };
1140        let chunk = Arc::try_unwrap(self.chunk).ok().unwrap();
1141        let data = chunk.into_inner().unwrap();
1142        Ok(ArrowColumnChunk { data, close })
1143    }
1144
1145    /// Returns the estimated total memory usage by the writer.
1146    ///
1147    /// This  [`Self::get_estimated_total_bytes`] this is an estimate
1148    /// of the current memory usage and not it's anticipated encoded size.
1149    ///
1150    /// This includes:
1151    /// 1. Data buffered in encoded form
1152    /// 2. Data buffered in un-encoded form (e.g. `usize` dictionary keys)
1153    ///
1154    /// This value should be greater than or equal to [`Self::get_estimated_total_bytes`]
1155    pub fn memory_size(&self) -> usize {
1156        match &self.writer {
1157            ArrowColumnWriterImpl::ByteArray(c) => c.memory_size(),
1158            ArrowColumnWriterImpl::Column(c) => c.memory_size(),
1159        }
1160    }
1161
1162    /// Returns the estimated total encoded bytes for this column writer.
1163    ///
1164    /// This includes:
1165    /// 1. Data buffered in encoded form
1166    /// 2. An estimate of how large the data buffered in un-encoded form would be once encoded
1167    ///
1168    /// This value should be less than or equal to [`Self::memory_size`]
1169    pub fn get_estimated_total_bytes(&self) -> usize {
1170        match &self.writer {
1171            ArrowColumnWriterImpl::ByteArray(c) => c.get_estimated_total_bytes() as _,
1172            ArrowColumnWriterImpl::Column(c) => c.get_estimated_total_bytes() as _,
1173        }
1174    }
1175}
1176
1177/// Encodes [`RecordBatch`] to a parquet row group
1178///
1179/// Note: this structure is created by [`ArrowRowGroupWriterFactory`] internally used to
1180/// create [`ArrowRowGroupWriter`]s, but it is not exposed publicly.
1181///
1182/// See the example on [`ArrowColumnWriter`] for how to encode columns in parallel
1183#[derive(Debug)]
1184struct ArrowRowGroupWriter {
1185    writers: Vec<ArrowColumnWriter>,
1186    schema: SchemaRef,
1187    buffered_rows: usize,
1188}
1189
1190impl ArrowRowGroupWriter {
1191    fn new(writers: Vec<ArrowColumnWriter>, arrow: &SchemaRef) -> Self {
1192        Self {
1193            writers,
1194            schema: arrow.clone(),
1195            buffered_rows: 0,
1196        }
1197    }
1198
1199    fn write(&mut self, batch: &RecordBatch) -> Result<()> {
1200        self.buffered_rows += batch.num_rows();
1201        let mut writers = self.writers.iter_mut();
1202        for (field, column) in self.schema.fields().iter().zip(batch.columns()) {
1203            for leaf in compute_leaves(field.as_ref(), column)? {
1204                writers.next().unwrap().write(&leaf)?;
1205            }
1206        }
1207        Ok(())
1208    }
1209
1210    fn write_with_chunkers(
1211        &mut self,
1212        batch: &RecordBatch,
1213        chunkers: &mut [ContentDefinedChunker],
1214    ) -> Result<()> {
1215        self.buffered_rows += batch.num_rows();
1216        let mut writers = self.writers.iter_mut();
1217        let mut chunkers = chunkers.iter_mut();
1218        for (field, column) in self.schema.fields().iter().zip(batch.columns()) {
1219            for leaf in compute_leaves(field.as_ref(), column)? {
1220                writers
1221                    .next()
1222                    .unwrap()
1223                    .write_with_chunker(&leaf, chunkers.next().unwrap())?;
1224            }
1225        }
1226        Ok(())
1227    }
1228
1229    /// Returns the estimated total encoded bytes for this row group
1230    fn get_estimated_total_bytes(&self) -> usize {
1231        self.writers
1232            .iter()
1233            .map(|x| x.get_estimated_total_bytes())
1234            .sum()
1235    }
1236
1237    fn close(self) -> Result<Vec<ArrowColumnChunk>> {
1238        self.writers
1239            .into_iter()
1240            .map(|writer| writer.close())
1241            .collect()
1242    }
1243}
1244
1245/// Factory that creates new column writers for each row group in the Parquet file.
1246///
1247/// You can create this structure via an [`ArrowWriter::into_serialized_writer`].
1248/// See the example on [`ArrowColumnWriter`] for how to encode columns in parallel
1249#[derive(Debug)]
1250pub struct ArrowRowGroupWriterFactory {
1251    schema: SchemaDescPtr,
1252    arrow_schema: SchemaRef,
1253    props: WriterPropertiesPtr,
1254    page_store_factory: Arc<dyn PageStoreFactory>,
1255    #[cfg(feature = "encryption")]
1256    file_encryptor: Option<Arc<FileEncryptor>>,
1257}
1258
1259impl ArrowRowGroupWriterFactory {
1260    /// Create a new [`ArrowRowGroupWriterFactory`] for the provided file writer and Arrow schema
1261    pub fn new<W: Write + Send>(
1262        file_writer: &SerializedFileWriter<W>,
1263        arrow_schema: SchemaRef,
1264    ) -> Self {
1265        let schema = Arc::clone(file_writer.schema_descr_ptr());
1266        let props = Arc::clone(file_writer.properties());
1267        Self {
1268            schema,
1269            arrow_schema,
1270            props,
1271            page_store_factory: Arc::new(InMemoryPageStoreFactory),
1272            #[cfg(feature = "encryption")]
1273            file_encryptor: file_writer.file_encryptor(),
1274        }
1275    }
1276
1277    /// Set the [`PageStoreFactory`] used to allocate the buffer for each column
1278    /// chunk, e.g. to spill completed pages to a temp file or object storage
1279    /// instead of the heap. Defaults to [`InMemoryPageStoreFactory`].
1280    pub fn with_page_store_factory(
1281        mut self,
1282        page_store_factory: Arc<dyn PageStoreFactory>,
1283    ) -> Self {
1284        self.page_store_factory = page_store_factory;
1285        self
1286    }
1287
1288    fn create_row_group_writer(&self, row_group_index: usize) -> Result<ArrowRowGroupWriter> {
1289        let writers = self.create_column_writers(row_group_index)?;
1290        Ok(ArrowRowGroupWriter::new(writers, &self.arrow_schema))
1291    }
1292
1293    /// Create column writers for a new row group, with the given row group index
1294    pub fn create_column_writers(&self, row_group_index: usize) -> Result<Vec<ArrowColumnWriter>> {
1295        let mut writers = Vec::with_capacity(self.arrow_schema.fields.len());
1296        let mut leaves = self.schema.columns().iter();
1297        let column_factory = self.column_writer_factory(row_group_index);
1298        for field in &self.arrow_schema.fields {
1299            column_factory.get_arrow_column_writer(
1300                field.data_type(),
1301                &self.props,
1302                &mut leaves,
1303                &mut writers,
1304            )?;
1305        }
1306        Ok(writers)
1307    }
1308
1309    #[cfg(feature = "encryption")]
1310    fn column_writer_factory(&self, row_group_idx: usize) -> ArrowColumnWriterFactory {
1311        ArrowColumnWriterFactory::new()
1312            .with_page_store_factory(self.page_store_factory.clone())
1313            .with_file_encryptor(row_group_idx, self.file_encryptor.clone())
1314    }
1315
1316    #[cfg(not(feature = "encryption"))]
1317    fn column_writer_factory(&self, _row_group_idx: usize) -> ArrowColumnWriterFactory {
1318        ArrowColumnWriterFactory::new().with_page_store_factory(self.page_store_factory.clone())
1319    }
1320}
1321
1322/// Creates [`ArrowColumnWriter`] instances
1323struct ArrowColumnWriterFactory {
1324    /// Allocates the per-column-chunk [`PageStore`] backing each page writer.
1325    page_store_factory: Arc<dyn PageStoreFactory>,
1326    #[cfg(feature = "encryption")]
1327    row_group_index: usize,
1328    #[cfg(feature = "encryption")]
1329    file_encryptor: Option<Arc<FileEncryptor>>,
1330}
1331
1332impl ArrowColumnWriterFactory {
1333    pub fn new() -> Self {
1334        Self {
1335            page_store_factory: Arc::new(InMemoryPageStoreFactory),
1336            #[cfg(feature = "encryption")]
1337            row_group_index: 0,
1338            #[cfg(feature = "encryption")]
1339            file_encryptor: None,
1340        }
1341    }
1342
1343    /// Use `page_store_factory` to allocate the buffer for each column chunk.
1344    pub fn with_page_store_factory(
1345        mut self,
1346        page_store_factory: Arc<dyn PageStoreFactory>,
1347    ) -> Self {
1348        self.page_store_factory = page_store_factory;
1349        self
1350    }
1351
1352    #[cfg(feature = "encryption")]
1353    pub fn with_file_encryptor(
1354        mut self,
1355        row_group_index: usize,
1356        file_encryptor: Option<Arc<FileEncryptor>>,
1357    ) -> Self {
1358        self.row_group_index = row_group_index;
1359        self.file_encryptor = file_encryptor;
1360        self
1361    }
1362
1363    #[cfg(feature = "encryption")]
1364    fn create_page_writer(
1365        &self,
1366        column_descriptor: &ColumnDescPtr,
1367        column_index: usize,
1368    ) -> Result<Box<ArrowPageWriter>> {
1369        let column_path = column_descriptor.path().string();
1370        let page_encryptor = PageEncryptor::create_if_column_encrypted(
1371            &self.file_encryptor,
1372            self.row_group_index,
1373            column_index,
1374            &column_path,
1375        )?;
1376        let args = PageStoreArgs::new(column_index, column_descriptor);
1377        let store = self.page_store_factory.create(&args)?;
1378        Ok(Box::new(
1379            ArrowPageWriter::new(store).with_encryptor(page_encryptor),
1380        ))
1381    }
1382
1383    #[cfg(not(feature = "encryption"))]
1384    fn create_page_writer(
1385        &self,
1386        column_descriptor: &ColumnDescPtr,
1387        column_index: usize,
1388    ) -> Result<Box<ArrowPageWriter>> {
1389        let args = PageStoreArgs::new(column_index, column_descriptor);
1390        let store = self.page_store_factory.create(&args)?;
1391        Ok(Box::new(ArrowPageWriter::new(store)))
1392    }
1393
1394    /// Gets an [`ArrowColumnWriter`] for the given `data_type`, appending the
1395    /// output ColumnDesc to `leaves` and the column writers to `out`
1396    fn get_arrow_column_writer(
1397        &self,
1398        data_type: &ArrowDataType,
1399        props: &WriterPropertiesPtr,
1400        leaves: &mut Iter<'_, ColumnDescPtr>,
1401        out: &mut Vec<ArrowColumnWriter>,
1402    ) -> Result<()> {
1403        // Instantiate writers for normal columns
1404        let col = |desc: &ColumnDescPtr| -> Result<ArrowColumnWriter> {
1405            let page_writer = self.create_page_writer(desc, out.len())?;
1406            let chunk = page_writer.buffer.clone();
1407            let writer = get_column_writer(desc.clone(), props.clone(), page_writer);
1408            Ok(ArrowColumnWriter {
1409                chunk,
1410                writer: ArrowColumnWriterImpl::Column(writer),
1411            })
1412        };
1413
1414        // Instantiate writers for byte arrays (e.g. Utf8,  Binary, etc)
1415        let bytes = |desc: &ColumnDescPtr| -> Result<ArrowColumnWriter> {
1416            let page_writer = self.create_page_writer(desc, out.len())?;
1417            let chunk = page_writer.buffer.clone();
1418            let writer = GenericColumnWriter::new(desc.clone(), props.clone(), page_writer);
1419            Ok(ArrowColumnWriter {
1420                chunk,
1421                writer: ArrowColumnWriterImpl::ByteArray(writer),
1422            })
1423        };
1424
1425        match data_type {
1426            _ if data_type.is_primitive() => out.push(col(leaves.next().unwrap())?),
1427            ArrowDataType::FixedSizeBinary(_) | ArrowDataType::Boolean | ArrowDataType::Null => {
1428                out.push(col(leaves.next().unwrap())?)
1429            }
1430            ArrowDataType::LargeBinary
1431            | ArrowDataType::Binary
1432            | ArrowDataType::Utf8
1433            | ArrowDataType::LargeUtf8
1434            | ArrowDataType::BinaryView
1435            | ArrowDataType::Utf8View => out.push(bytes(leaves.next().unwrap())?),
1436            ArrowDataType::List(f)
1437            | ArrowDataType::LargeList(f)
1438            | ArrowDataType::FixedSizeList(f, _)
1439            | ArrowDataType::ListView(f)
1440            | ArrowDataType::LargeListView(f) => {
1441                self.get_arrow_column_writer(f.data_type(), props, leaves, out)?
1442            }
1443            ArrowDataType::Struct(fields) => {
1444                for field in fields {
1445                    self.get_arrow_column_writer(field.data_type(), props, leaves, out)?
1446                }
1447            }
1448            ArrowDataType::Map(f, _) => match f.data_type() {
1449                ArrowDataType::Struct(f) => {
1450                    self.get_arrow_column_writer(f[0].data_type(), props, leaves, out)?;
1451                    self.get_arrow_column_writer(f[1].data_type(), props, leaves, out)?
1452                }
1453                _ => unreachable!("invalid map type"),
1454            },
1455            ArrowDataType::Dictionary(_, value_type) => match value_type.as_ref() {
1456                ArrowDataType::Utf8
1457                | ArrowDataType::LargeUtf8
1458                | ArrowDataType::Binary
1459                | ArrowDataType::LargeBinary => out.push(bytes(leaves.next().unwrap())?),
1460                ArrowDataType::Utf8View | ArrowDataType::BinaryView => {
1461                    out.push(bytes(leaves.next().unwrap())?)
1462                }
1463                ArrowDataType::FixedSizeBinary(_) => out.push(bytes(leaves.next().unwrap())?),
1464                _ => out.push(col(leaves.next().unwrap())?),
1465            },
1466            ArrowDataType::RunEndEncoded(_, value_field) => {
1467                self.get_arrow_column_writer(value_field.data_type(), props, leaves, out)?
1468            }
1469            _ => {
1470                return Err(ParquetError::NYI(format!(
1471                    "Attempting to write an Arrow type {data_type} to parquet that is not yet implemented"
1472                )));
1473            }
1474        }
1475        Ok(())
1476    }
1477}
1478
1479fn write_leaf(
1480    writer: &mut ColumnWriter<'_>,
1481    column: &dyn arrow_array::Array,
1482    levels: &ArrayLevels,
1483) -> Result<usize> {
1484    let indices = levels.non_null_indices();
1485
1486    match writer {
1487        // Note: this should match the contents of arrow_to_parquet_type
1488        ColumnWriter::Int32ColumnWriter(typed) => {
1489            match column.data_type() {
1490                ArrowDataType::Null => {
1491                    let array = Int32Array::new_null(column.len());
1492                    write_primitive(typed, array.values(), levels)
1493                }
1494                ArrowDataType::Int8 => {
1495                    let array: Int32Array = column.as_primitive::<Int8Type>().unary(|x| x as i32);
1496                    write_primitive(typed, array.values(), levels)
1497                }
1498                ArrowDataType::Int16 => {
1499                    let array: Int32Array = column.as_primitive::<Int16Type>().unary(|x| x as i32);
1500                    write_primitive(typed, array.values(), levels)
1501                }
1502                ArrowDataType::Int32 => {
1503                    write_primitive(typed, column.as_primitive::<Int32Type>().values(), levels)
1504                }
1505                ArrowDataType::UInt8 => {
1506                    let array: Int32Array = column.as_primitive::<UInt8Type>().unary(|x| x as i32);
1507                    write_primitive(typed, array.values(), levels)
1508                }
1509                ArrowDataType::UInt16 => {
1510                    let array: Int32Array = column.as_primitive::<UInt16Type>().unary(|x| x as i32);
1511                    write_primitive(typed, array.values(), levels)
1512                }
1513                ArrowDataType::UInt32 => {
1514                    // follow C++ implementation and use overflow/reinterpret cast from  u32 to i32 which will map
1515                    // `(i32::MAX as u32)..u32::MAX` to `i32::MIN..0`
1516                    let array = column.as_primitive::<UInt32Type>();
1517                    write_primitive(typed, array.values().inner().typed_data(), levels)
1518                }
1519                ArrowDataType::Date32 => {
1520                    let array = column.as_primitive::<Date32Type>();
1521                    write_primitive(typed, array.values(), levels)
1522                }
1523                ArrowDataType::Time32(TimeUnit::Second) => {
1524                    let array = column.as_primitive::<Time32SecondType>();
1525                    write_primitive(typed, array.values(), levels)
1526                }
1527                ArrowDataType::Time32(TimeUnit::Millisecond) => {
1528                    let array = column.as_primitive::<Time32MillisecondType>();
1529                    write_primitive(typed, array.values(), levels)
1530                }
1531                ArrowDataType::Date64 => {
1532                    // If the column is a Date64, we truncate it
1533                    let array: Int32Array = column
1534                        .as_primitive::<Date64Type>()
1535                        .unary(|x| (x / 86_400_000) as _);
1536
1537                    write_primitive(typed, array.values(), levels)
1538                }
1539                ArrowDataType::Decimal32(_, _) => {
1540                    let array = column
1541                        .as_primitive::<Decimal32Type>()
1542                        .unary::<_, Int32Type>(|v| v);
1543                    write_primitive(typed, array.values(), levels)
1544                }
1545                ArrowDataType::Decimal64(_, _) => {
1546                    // use the int32 to represent the decimal with low precision
1547                    let array = column
1548                        .as_primitive::<Decimal64Type>()
1549                        .unary::<_, Int32Type>(|v| v as i32);
1550                    write_primitive(typed, array.values(), levels)
1551                }
1552                ArrowDataType::Decimal128(_, _) => {
1553                    // use the int32 to represent the decimal with low precision
1554                    let array = column
1555                        .as_primitive::<Decimal128Type>()
1556                        .unary::<_, Int32Type>(|v| v as i32);
1557                    write_primitive(typed, array.values(), levels)
1558                }
1559                ArrowDataType::Decimal256(_, _) => {
1560                    // use the int32 to represent the decimal with low precision
1561                    let array = column
1562                        .as_primitive::<Decimal256Type>()
1563                        .unary::<_, Int32Type>(|v| v.as_i128() as i32);
1564                    write_primitive(typed, array.values(), levels)
1565                }
1566                d => Err(ParquetError::General(format!("Cannot coerce {d} to I32"))),
1567            }
1568        }
1569        ColumnWriter::BoolColumnWriter(typed) => {
1570            let array = column.as_boolean();
1571            let values = get_bool_array_slice(array, indices.iter().copied());
1572            typed.write_batch_internal(
1573                values.as_slice(),
1574                None,
1575                levels.def_level_data().as_ref(),
1576                levels.rep_level_data().as_ref(),
1577                None,
1578                None,
1579                None,
1580            )
1581        }
1582        ColumnWriter::Int64ColumnWriter(typed) => {
1583            match column.data_type() {
1584                ArrowDataType::Date64 => {
1585                    let array = column
1586                        .as_primitive::<Date64Type>()
1587                        .reinterpret_cast::<Int64Type>();
1588
1589                    write_primitive(typed, array.values(), levels)
1590                }
1591                ArrowDataType::Int64 => {
1592                    let array = column.as_primitive::<Int64Type>();
1593                    write_primitive(typed, array.values(), levels)
1594                }
1595                ArrowDataType::UInt64 => {
1596                    let values = column.as_primitive::<UInt64Type>().values();
1597                    // follow C++ implementation and use overflow/reinterpret cast from  u64 to i64 which will map
1598                    // `(i64::MAX as u64)..u64::MAX` to `i64::MIN..0`
1599                    let array = values.inner().typed_data::<i64>();
1600                    write_primitive(typed, array, levels)
1601                }
1602                ArrowDataType::Time64(TimeUnit::Microsecond) => {
1603                    let array = column.as_primitive::<Time64MicrosecondType>();
1604                    write_primitive(typed, array.values(), levels)
1605                }
1606                ArrowDataType::Time64(TimeUnit::Nanosecond) => {
1607                    let array = column.as_primitive::<Time64NanosecondType>();
1608                    write_primitive(typed, array.values(), levels)
1609                }
1610                ArrowDataType::Timestamp(unit, _) => match unit {
1611                    TimeUnit::Second => {
1612                        let array = column.as_primitive::<TimestampSecondType>();
1613                        write_primitive(typed, array.values(), levels)
1614                    }
1615                    TimeUnit::Millisecond => {
1616                        let array = column.as_primitive::<TimestampMillisecondType>();
1617                        write_primitive(typed, array.values(), levels)
1618                    }
1619                    TimeUnit::Microsecond => {
1620                        let array = column.as_primitive::<TimestampMicrosecondType>();
1621                        write_primitive(typed, array.values(), levels)
1622                    }
1623                    TimeUnit::Nanosecond => {
1624                        let array = column.as_primitive::<TimestampNanosecondType>();
1625                        write_primitive(typed, array.values(), levels)
1626                    }
1627                },
1628                ArrowDataType::Duration(unit) => match unit {
1629                    TimeUnit::Second => {
1630                        let array = column.as_primitive::<DurationSecondType>();
1631                        write_primitive(typed, array.values(), levels)
1632                    }
1633                    TimeUnit::Millisecond => {
1634                        let array = column.as_primitive::<DurationMillisecondType>();
1635                        write_primitive(typed, array.values(), levels)
1636                    }
1637                    TimeUnit::Microsecond => {
1638                        let array = column.as_primitive::<DurationMicrosecondType>();
1639                        write_primitive(typed, array.values(), levels)
1640                    }
1641                    TimeUnit::Nanosecond => {
1642                        let array = column.as_primitive::<DurationNanosecondType>();
1643                        write_primitive(typed, array.values(), levels)
1644                    }
1645                },
1646                ArrowDataType::Decimal64(_, _) => {
1647                    let array = column
1648                        .as_primitive::<Decimal64Type>()
1649                        .reinterpret_cast::<Int64Type>();
1650                    write_primitive(typed, array.values(), levels)
1651                }
1652                ArrowDataType::Decimal128(_, _) => {
1653                    // use the int64 to represent the decimal with low precision
1654                    let array = column
1655                        .as_primitive::<Decimal128Type>()
1656                        .unary::<_, Int64Type>(|v| v as i64);
1657                    write_primitive(typed, array.values(), levels)
1658                }
1659                ArrowDataType::Decimal256(_, _) => {
1660                    // use the int64 to represent the decimal with low precision
1661                    let array = column
1662                        .as_primitive::<Decimal256Type>()
1663                        .unary::<_, Int64Type>(|v| v.as_i128() as i64);
1664                    write_primitive(typed, array.values(), levels)
1665                }
1666                d => Err(ParquetError::General(format!("Cannot coerce {d} to I64"))),
1667            }
1668        }
1669        ColumnWriter::Int96ColumnWriter(_typed) => {
1670            unreachable!("Currently unreachable because data type not supported")
1671        }
1672        ColumnWriter::FloatColumnWriter(typed) => {
1673            let array = column.as_primitive::<Float32Type>();
1674            write_primitive(typed, array.values(), levels)
1675        }
1676        ColumnWriter::DoubleColumnWriter(typed) => {
1677            let array = column.as_primitive::<Float64Type>();
1678            write_primitive(typed, array.values(), levels)
1679        }
1680        ColumnWriter::ByteArrayColumnWriter(_) => {
1681            unreachable!("should use ByteArrayWriter")
1682        }
1683        ColumnWriter::FixedLenByteArrayColumnWriter(typed) => {
1684            let bytes = match column.data_type() {
1685                ArrowDataType::Interval(interval_unit) => match interval_unit {
1686                    IntervalUnit::YearMonth => {
1687                        let array = column.as_primitive::<IntervalYearMonthType>();
1688                        get_interval_ym_array_slice(array, indices.iter().copied())
1689                    }
1690                    IntervalUnit::DayTime => {
1691                        let array = column.as_primitive::<IntervalDayTimeType>();
1692                        get_interval_dt_array_slice(array, indices.iter().copied())
1693                    }
1694                    _ => {
1695                        return Err(ParquetError::NYI(format!(
1696                            "Attempting to write an Arrow interval type {interval_unit:?} to parquet that is not yet implemented"
1697                        )));
1698                    }
1699                },
1700                ArrowDataType::FixedSizeBinary(_) => {
1701                    let array = column.as_fixed_size_binary();
1702                    get_fsb_array_slice(array, indices.iter().copied())
1703                }
1704                ArrowDataType::Decimal32(_, _) => {
1705                    let array = column.as_primitive::<Decimal32Type>();
1706                    get_decimal_array_slice(array, indices.iter().copied())
1707                }
1708                ArrowDataType::Decimal64(_, _) => {
1709                    let array = column.as_primitive::<Decimal64Type>();
1710                    get_decimal_array_slice(array, indices.iter().copied())
1711                }
1712                ArrowDataType::Decimal128(_, _) => {
1713                    let array = column.as_primitive::<Decimal128Type>();
1714                    get_decimal_array_slice(array, indices.iter().copied())
1715                }
1716                ArrowDataType::Decimal256(_, _) => {
1717                    let array = column.as_primitive::<Decimal256Type>();
1718                    get_decimal_array_slice(array, indices.iter().copied())
1719                }
1720                ArrowDataType::Float16 => {
1721                    let array = column.as_primitive::<Float16Type>();
1722                    get_float_16_array_slice(array, indices.iter().copied())
1723                }
1724                _ => {
1725                    return Err(ParquetError::NYI(
1726                        "Attempting to write an Arrow type that is not yet implemented".to_string(),
1727                    ));
1728                }
1729            };
1730            typed.write_batch_internal(
1731                bytes.as_slice(),
1732                None,
1733                levels.def_level_data().as_ref(),
1734                levels.rep_level_data().as_ref(),
1735                None,
1736                None,
1737                None,
1738            )
1739        }
1740    }
1741}
1742
1743fn write_primitive<E: ColumnValueEncoder>(
1744    writer: &mut GenericColumnWriter<E>,
1745    values: &E::Values,
1746    levels: &ArrayLevels,
1747) -> Result<usize> {
1748    writer.write_batch_internal(
1749        values,
1750        Some(levels.non_null_indices()),
1751        levels.def_level_data().as_ref(),
1752        levels.rep_level_data().as_ref(),
1753        None,
1754        None,
1755        None,
1756    )
1757}
1758
1759fn get_bool_array_slice(
1760    array: &arrow_array::BooleanArray,
1761    indices: impl ExactSizeIterator<Item = usize>,
1762) -> Vec<bool> {
1763    let mut values = Vec::with_capacity(indices.len());
1764    for i in indices {
1765        values.push(array.value(i))
1766    }
1767    values
1768}
1769
1770/// Returns 12-byte values representing 3 values of months, days and milliseconds (4-bytes each).
1771/// An Arrow YearMonth interval only stores months, thus only the first 4 bytes are populated.
1772fn get_interval_ym_array_slice(
1773    array: &arrow_array::IntervalYearMonthArray,
1774    indices: impl ExactSizeIterator<Item = usize>,
1775) -> Vec<FixedLenByteArray> {
1776    chunk_array_slice(12, indices, move |i, chunk| {
1777        let value = array.value(i);
1778        chunk[0..4].copy_from_slice(&value.to_le_bytes());
1779    })
1780}
1781
1782/// Returns 12-byte values representing 3 values of months, days and milliseconds (4-bytes each).
1783/// An Arrow DayTime interval only stores days and millis, thus the first 4 bytes are not populated.
1784fn get_interval_dt_array_slice(
1785    array: &arrow_array::IntervalDayTimeArray,
1786    indices: impl ExactSizeIterator<Item = usize>,
1787) -> Vec<FixedLenByteArray> {
1788    chunk_array_slice(12, indices, move |i, chunk| {
1789        let value = array.value(i);
1790        chunk[4..8].copy_from_slice(&value.days.to_le_bytes());
1791        chunk[8..12].copy_from_slice(&value.milliseconds.to_le_bytes());
1792    })
1793}
1794
1795trait NativeDecimalType: DecimalType {
1796    type NativeBytes: AsRef<[u8]>;
1797
1798    fn to_be_bytes(value: Self::Native) -> Self::NativeBytes;
1799}
1800impl NativeDecimalType for Decimal32Type {
1801    type NativeBytes = [u8; Self::BYTE_LENGTH];
1802
1803    fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
1804        value.to_be_bytes()
1805    }
1806}
1807impl NativeDecimalType for Decimal64Type {
1808    type NativeBytes = [u8; Self::BYTE_LENGTH];
1809
1810    fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
1811        value.to_be_bytes()
1812    }
1813}
1814impl NativeDecimalType for Decimal128Type {
1815    type NativeBytes = [u8; Self::BYTE_LENGTH];
1816
1817    fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
1818        value.to_be_bytes()
1819    }
1820}
1821impl NativeDecimalType for Decimal256Type {
1822    type NativeBytes = [u8; Self::BYTE_LENGTH];
1823
1824    fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
1825        value.to_be_bytes()
1826    }
1827}
1828
1829fn get_decimal_array_slice<T: NativeDecimalType>(
1830    array: &PrimitiveArray<T>,
1831    indices: impl ExactSizeIterator<Item = usize>,
1832) -> Vec<FixedLenByteArray> {
1833    let chunk_size = decimal_length_from_precision(array.precision());
1834    assert!(chunk_size <= T::BYTE_LENGTH);
1835
1836    if chunk_size == T::BYTE_LENGTH {
1837        // Special-case that allows inlining memcpy.
1838        chunk_array_slice(chunk_size, indices, move |i, chunk| {
1839            let as_be_bytes = T::to_be_bytes(array.value(i));
1840            chunk.copy_from_slice(as_be_bytes.as_ref());
1841        })
1842    } else {
1843        chunk_array_slice(chunk_size, indices, move |i, chunk| {
1844            let as_be_bytes = T::to_be_bytes(array.value(i));
1845            let resized_value = &as_be_bytes.as_ref()[(T::BYTE_LENGTH - chunk.len())..];
1846            chunk.copy_from_slice(resized_value);
1847        })
1848    }
1849}
1850
1851fn get_float_16_array_slice(
1852    array: &arrow_array::Float16Array,
1853    indices: impl ExactSizeIterator<Item = usize>,
1854) -> Vec<FixedLenByteArray> {
1855    chunk_array_slice(2, indices, move |i, chunk| {
1856        let value = array.value(i).to_le_bytes();
1857        chunk.copy_from_slice(&value);
1858    })
1859}
1860
1861fn get_fsb_array_slice(
1862    array: &arrow_array::FixedSizeBinaryArray,
1863    indices: impl ExactSizeIterator<Item = usize>,
1864) -> Vec<FixedLenByteArray> {
1865    chunk_array_slice(array.value_size(), indices, move |i, chunk| {
1866        let value = array.value(i);
1867        chunk.copy_from_slice(value);
1868    })
1869}
1870
1871#[inline]
1872fn chunk_array_slice(
1873    chunk_size: usize,
1874    indices: impl ExactSizeIterator<Item = usize>,
1875    writer: impl Fn(usize, &mut [u8]),
1876) -> Vec<FixedLenByteArray> {
1877    let capacity = indices.len() * chunk_size;
1878    // TODO: This could be done with Vec::spare_capacity_mut,
1879    //       but [MaybeUninit]::write_copy_of_slice is gated behind MSRV 1.93
1880    let mut arena = vec![0; capacity];
1881    for (i, chunk) in indices.zip(arena.chunks_exact_mut(chunk_size)) {
1882        writer(i, chunk);
1883    }
1884    chunk_contiguous_vec(arena, chunk_size)
1885}
1886
1887fn chunk_contiguous_vec(arena: Vec<u8>, chunk_size: usize) -> Vec<FixedLenByteArray> {
1888    let mut values = Vec::with_capacity(arena.len() / chunk_size);
1889    let mut arena = Bytes::from(arena);
1890    while arena.len() >= chunk_size {
1891        let slice = arena.split_to(chunk_size);
1892        values.push(FixedLenByteArray::from(ByteArray::from(slice)));
1893    }
1894    values
1895}
1896
1897#[cfg(test)]
1898mod tests {
1899    use super::*;
1900    use std::cmp::Ordering;
1901    use std::collections::HashMap;
1902
1903    use std::fs::File;
1904
1905    use crate::arrow::arrow_reader::{ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder};
1906    use crate::arrow::{ARROW_SCHEMA_META_KEY, PARQUET_FIELD_ID_META_KEY};
1907    use crate::column::page::{Page, PageReader};
1908    use crate::file::metadata::thrift::PageHeader;
1909    use crate::file::page_index::column_index::ColumnIndexMetaData;
1910    use crate::file::reader::SerializedPageReader;
1911    use crate::parquet_thrift::{ReadThrift, ThriftSliceInputProtocol};
1912    use crate::schema::types::ColumnPath;
1913    use arrow::datatypes::ToByteSlice;
1914    use arrow::datatypes::{DataType, Schema};
1915    use arrow::error::Result as ArrowResult;
1916    use arrow::util::data_gen::create_random_array;
1917    use arrow::util::pretty::pretty_format_batches;
1918    use arrow::{array::*, buffer::Buffer};
1919    use arrow_buffer::{IntervalDayTime, IntervalMonthDayNano, NullBuffer, OffsetBuffer, i256};
1920    use arrow_schema::Fields;
1921    use half::f16;
1922    use num_traits::{FromPrimitive, ToPrimitive};
1923    use tempfile::tempfile;
1924
1925    use crate::basic::Encoding;
1926    use crate::data_type::AsBytes;
1927    use crate::file::metadata::{ColumnChunkMetaData, ParquetMetaData, ParquetMetaDataReader};
1928    use crate::file::properties::{
1929        BloomFilterPosition, EnabledStatistics, ReaderProperties, WriterVersion,
1930    };
1931    use crate::file::serialized_reader::ReadOptionsBuilder;
1932    use crate::file::{
1933        reader::{FileReader, SerializedFileReader},
1934        statistics::Statistics,
1935    };
1936
1937    /// A [`PageStore`] that allocates *sparse, non-contiguous* handles and keeps
1938    /// blobs in a `HashMap` — nothing like the default `Vec<Bytes>`. Used to
1939    /// prove the writer relies only on the opaque-handle contract and never on
1940    /// handles being dense `Vec` indices. Records how many blobs were stored.
1941    #[derive(Debug, Default)]
1942    struct RecordingPageStore {
1943        next: u64,
1944        blobs: HashMap<u64, Bytes>,
1945        puts: Arc<std::sync::atomic::AtomicUsize>,
1946    }
1947
1948    impl PageStore for RecordingPageStore {
1949        fn put(&mut self, value: Bytes) -> Result<PageKey> {
1950            // Deliberately non-sequential, never-zero handles.
1951            let id = 100 + self.next * 7;
1952            self.next += 1;
1953            self.puts.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1954            self.blobs.insert(id, value);
1955            Ok(PageKey::new(id))
1956        }
1957
1958        fn take(&mut self, key: PageKey) -> Result<Bytes> {
1959            self.blobs
1960                .remove(&key.get())
1961                .ok_or_else(|| ParquetError::General(format!("missing key {}", key.get())))
1962        }
1963    }
1964
1965    #[derive(Debug)]
1966    struct RecordingPageStoreFactory {
1967        puts: Arc<std::sync::atomic::AtomicUsize>,
1968    }
1969
1970    impl PageStoreFactory for RecordingPageStoreFactory {
1971        fn create(&self, _args: &PageStoreArgs<'_>) -> Result<Box<dyn PageStore>> {
1972            Ok(Box::new(RecordingPageStore {
1973                puts: self.puts.clone(),
1974                ..Default::default()
1975            }))
1976        }
1977    }
1978
1979    /// A custom [`PageStore`] must produce byte-identical files to the in-memory
1980    /// default, across dictionary and non-dictionary columns and multiple row
1981    /// groups (so multiple store instances are exercised).
1982    #[test]
1983    fn custom_page_store_is_byte_identical_to_default() {
1984        let schema = Arc::new(Schema::new(vec![
1985            Field::new("i", DataType::Int32, true),
1986            // A low-cardinality string column to exercise the dictionary path.
1987            Field::new("s", DataType::Utf8, true),
1988        ]));
1989        let i = Int32Array::from(vec![Some(1), None, Some(3), Some(4), Some(5), Some(6)]);
1990        let s = StringArray::from(vec![
1991            Some("a"),
1992            Some("bb"),
1993            Some("a"),
1994            None,
1995            Some("bb"),
1996            Some("ccc"),
1997        ]);
1998        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(i), Arc::new(s)]).unwrap();
1999
2000        // Small row groups so multiple column chunks (hence multiple store
2001        // instances) are produced.
2002        let props = WriterProperties::builder()
2003            .set_max_row_group_row_count(Some(3))
2004            .build();
2005
2006        let write = |factory: Option<Arc<dyn PageStoreFactory>>| {
2007            let mut buffer = Vec::new();
2008            let mut opts = ArrowWriterOptions::new().with_properties(props.clone());
2009            if let Some(factory) = factory {
2010                opts = opts.with_page_store_factory(factory);
2011            }
2012            let mut writer =
2013                ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap();
2014            writer.write(&batch).unwrap();
2015            writer.close().unwrap();
2016            buffer
2017        };
2018
2019        let default_bytes = write(None);
2020
2021        let puts = Arc::new(std::sync::atomic::AtomicUsize::new(0));
2022        let custom_bytes = write(Some(Arc::new(RecordingPageStoreFactory {
2023            puts: puts.clone(),
2024        })));
2025
2026        assert!(
2027            puts.load(std::sync::atomic::Ordering::Relaxed) > 0,
2028            "custom PageStore was never written to"
2029        );
2030        assert_eq!(
2031            default_bytes, custom_bytes,
2032            "a custom PageStore must produce byte-identical output to the default"
2033        );
2034    }
2035
2036    /// A dictionary-encoded column written through the deferred-ordering Arrow
2037    /// path must round-trip correctly even with the offset index disabled, when
2038    /// only the chunk-level dictionary/data page offsets are rewritten (there is
2039    /// no offset index to rebuild). Spans multiple data pages so the
2040    /// dictionary-first reordering is exercised.
2041    #[test]
2042    fn dictionary_column_round_trips_with_offset_index_disabled() {
2043        let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int32, true)]));
2044
2045        // Low cardinality so the column stays dictionary-encoded; enough rows to
2046        // span several data pages within a single row group.
2047        let values: Vec<Option<i32>> = (0..50_000).map(|i| Some(i % 8)).collect();
2048        let array = Int32Array::from(values.clone());
2049        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
2050
2051        let props = WriterProperties::builder()
2052            .set_offset_index_disabled(true)
2053            .set_data_page_row_count_limit(4096)
2054            .build();
2055        let opts = ArrowWriterOptions::new().with_properties(props);
2056
2057        let mut buffer = Vec::new();
2058        let mut writer =
2059            ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap();
2060        writer.write(&batch).unwrap();
2061        writer.close().unwrap();
2062
2063        let reader = ParquetRecordBatchReader::try_new(Bytes::from(buffer), values.len()).unwrap();
2064        let read: Vec<RecordBatch> = reader.collect::<ArrowResult<_>>().unwrap();
2065        let read_values: Vec<Option<i32>> = read
2066            .iter()
2067            .flat_map(|b| b.column(0).as_primitive::<Int32Type>().iter())
2068            .collect();
2069        assert_eq!(read_values, values);
2070    }
2071
2072    /// The dictionary page is routed through the [`PageStore`] like any other
2073    /// page rather than held resident in memory, so a dictionary column chunk's
2074    /// *entire* serialized size — dictionary page included — passes through the
2075    /// store.
2076    #[test]
2077    fn dictionary_page_is_routed_through_the_store() {
2078        /// A store that sums the bytes handed to `put`.
2079        #[derive(Debug, Default)]
2080        struct SizeRecordingPageStore {
2081            blobs: Vec<Bytes>,
2082            bytes_put: Arc<std::sync::atomic::AtomicUsize>,
2083        }
2084        impl PageStore for SizeRecordingPageStore {
2085            fn put(&mut self, value: Bytes) -> Result<PageKey> {
2086                self.bytes_put
2087                    .fetch_add(value.len(), std::sync::atomic::Ordering::Relaxed);
2088                let key = PageKey::new(self.blobs.len() as u64);
2089                self.blobs.push(value);
2090                Ok(key)
2091            }
2092            fn take(&mut self, key: PageKey) -> Result<Bytes> {
2093                Ok(std::mem::take(&mut self.blobs[key.get() as usize]))
2094            }
2095        }
2096        #[derive(Debug)]
2097        struct Factory {
2098            bytes_put: Arc<std::sync::atomic::AtomicUsize>,
2099        }
2100        impl PageStoreFactory for Factory {
2101            fn create(&self, _args: &PageStoreArgs<'_>) -> Result<Box<dyn PageStore>> {
2102                Ok(Box::new(SizeRecordingPageStore {
2103                    bytes_put: self.bytes_put.clone(),
2104                    ..Default::default()
2105                }))
2106            }
2107        }
2108
2109        let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)]));
2110        // Low cardinality keeps the column dictionary-encoded with a real,
2111        // non-empty dictionary page.
2112        let values: Vec<&str> = (0..2048)
2113            .map(|i| ["alpha", "beta", "gamma", "delta"][i % 4])
2114            .collect();
2115        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(StringArray::from(values))])
2116            .unwrap();
2117
2118        let bytes_put = Arc::new(std::sync::atomic::AtomicUsize::new(0));
2119        let opts = ArrowWriterOptions::new().with_page_store_factory(Arc::new(Factory {
2120            bytes_put: bytes_put.clone(),
2121        }));
2122
2123        // A single batch / single column means exactly one row group and one
2124        // store instance, so the bytes it saw map to one column chunk.
2125        let mut buffer = Vec::new();
2126        let mut writer =
2127            ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap();
2128        writer.write(&batch).unwrap();
2129        writer.close().unwrap();
2130
2131        let reader = SerializedFileReader::new(Bytes::from(buffer)).unwrap();
2132        let column = reader.metadata().row_group(0).column(0);
2133        assert!(
2134            column.dictionary_page_offset().is_some(),
2135            "expected the column to be dictionary-encoded"
2136        );
2137
2138        // The bytes the store was handed must account for the whole chunk,
2139        // dictionary page included. Holding the dictionary page apart from the
2140        // store would make this fall short by the dictionary page's size.
2141        assert_eq!(
2142            bytes_put.load(std::sync::atomic::Ordering::Relaxed) as i64,
2143            column.compressed_size(),
2144            "the dictionary page must pass through the store like any other page"
2145        );
2146    }
2147
2148    #[test]
2149    fn arrow_writer() {
2150        // define schema
2151        let schema = Schema::new(vec![
2152            Field::new("a", DataType::Int32, false),
2153            Field::new("b", DataType::Int32, true),
2154        ]);
2155
2156        // create some data
2157        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2158        let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2159
2160        // build a record batch
2161        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]).unwrap();
2162
2163        roundtrip(batch, Some(SMALL_SIZE / 2));
2164    }
2165
2166    fn get_bytes_after_close(schema: SchemaRef, expected_batch: &RecordBatch) -> Vec<u8> {
2167        let mut buffer = vec![];
2168
2169        let mut writer = ArrowWriter::try_new(&mut buffer, schema, None).unwrap();
2170        writer.write(expected_batch).unwrap();
2171        writer.close().unwrap();
2172
2173        buffer
2174    }
2175
2176    fn get_bytes_by_into_inner(schema: SchemaRef, expected_batch: &RecordBatch) -> Vec<u8> {
2177        let mut writer = ArrowWriter::try_new(Vec::new(), schema, None).unwrap();
2178        writer.write(expected_batch).unwrap();
2179        writer.into_inner().unwrap()
2180    }
2181
2182    #[test]
2183    fn roundtrip_bytes() {
2184        // define schema
2185        let schema = Arc::new(Schema::new(vec![
2186            Field::new("a", DataType::Int32, false),
2187            Field::new("b", DataType::Int32, true),
2188        ]));
2189
2190        // create some data
2191        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2192        let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2193
2194        // build a record batch
2195        let expected_batch =
2196            RecordBatch::try_new(schema.clone(), vec![Arc::new(a), Arc::new(b)]).unwrap();
2197
2198        for buffer in [
2199            get_bytes_after_close(schema.clone(), &expected_batch),
2200            get_bytes_by_into_inner(schema, &expected_batch),
2201        ] {
2202            let cursor = Bytes::from(buffer);
2203            let mut record_batch_reader = ParquetRecordBatchReader::try_new(cursor, 1024).unwrap();
2204
2205            let actual_batch = record_batch_reader
2206                .next()
2207                .expect("No batch found")
2208                .expect("Unable to get batch");
2209
2210            assert_eq!(expected_batch.schema(), actual_batch.schema());
2211            assert_eq!(expected_batch.num_columns(), actual_batch.num_columns());
2212            assert_eq!(expected_batch.num_rows(), actual_batch.num_rows());
2213            for i in 0..expected_batch.num_columns() {
2214                let expected_data = expected_batch.column(i).to_data();
2215                let actual_data = actual_batch.column(i).to_data();
2216
2217                assert_eq!(expected_data, actual_data);
2218            }
2219        }
2220    }
2221
2222    #[test]
2223    fn arrow_writer_non_null() {
2224        // define schema
2225        let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
2226
2227        // create some data
2228        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2229
2230        // build a record batch
2231        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2232
2233        roundtrip(batch, Some(SMALL_SIZE / 2));
2234    }
2235
2236    #[test]
2237    fn arrow_writer_list() {
2238        // define schema
2239        let schema = Schema::new(vec![Field::new(
2240            "a",
2241            DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false))),
2242            true,
2243        )]);
2244
2245        // create some data
2246        let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
2247
2248        // Construct a buffer for value offsets, for the nested array:
2249        //  [[1], [2, 3], null, [4, 5, 6], [7, 8, 9, 10]]
2250        let a_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice());
2251
2252        // Construct a list array from the above two
2253        let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
2254            DataType::Int32,
2255            false,
2256        ))))
2257        .len(5)
2258        .add_buffer(a_value_offsets)
2259        .add_child_data(a_values.into_data())
2260        .null_bit_buffer(Some(Buffer::from([0b00011011])))
2261        .build()
2262        .unwrap();
2263        let a = ListArray::from(a_list_data);
2264
2265        // build a record batch
2266        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2267
2268        assert_eq!(batch.column(0).null_count(), 1);
2269
2270        // This test fails if the max row group size is less than the batch's length
2271        // see https://github.com/apache/arrow-rs/issues/518
2272        roundtrip(batch, None);
2273    }
2274
2275    #[test]
2276    fn arrow_writer_list_non_null() {
2277        // define schema
2278        let schema = Schema::new(vec![Field::new(
2279            "a",
2280            DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false))),
2281            false,
2282        )]);
2283
2284        // create some data
2285        let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
2286
2287        // Construct a buffer for value offsets, for the nested array:
2288        //  [[1], [2, 3], [], [4, 5, 6], [7, 8, 9, 10]]
2289        let a_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice());
2290
2291        // Construct a list array from the above two
2292        let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
2293            DataType::Int32,
2294            false,
2295        ))))
2296        .len(5)
2297        .add_buffer(a_value_offsets)
2298        .add_child_data(a_values.into_data())
2299        .build()
2300        .unwrap();
2301        let a = ListArray::from(a_list_data);
2302
2303        // build a record batch
2304        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2305
2306        // This test fails if the max row group size is less than the batch's length
2307        // see https://github.com/apache/arrow-rs/issues/518
2308        assert_eq!(batch.column(0).null_count(), 0);
2309
2310        roundtrip(batch, None);
2311    }
2312
2313    #[test]
2314    fn arrow_writer_list_view() {
2315        let list_field = Arc::new(Field::new_list_field(DataType::Int32, false));
2316        let schema = Schema::new(vec![Field::new(
2317            "a",
2318            DataType::ListView(list_field.clone()),
2319            true,
2320        )]);
2321
2322        //  [[1], [2, 3], null, [4, 5, 6], [7, 8, 9, 10]]
2323        let a = ListViewArray::new(
2324            list_field,
2325            vec![0, 1, 0, 3, 6].into(),
2326            vec![1, 2, 0, 3, 4].into(),
2327            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
2328            Some(vec![true, true, false, true, true].into()),
2329        );
2330
2331        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2332
2333        assert_eq!(batch.column(0).null_count(), 1);
2334
2335        roundtrip(batch, None);
2336    }
2337
2338    #[test]
2339    fn arrow_writer_list_view_non_null() {
2340        let list_field = Arc::new(Field::new_list_field(DataType::Int32, false));
2341        let schema = Schema::new(vec![Field::new(
2342            "a",
2343            DataType::ListView(list_field.clone()),
2344            false,
2345        )]);
2346
2347        //  [[1], [2, 3], [], [4, 5, 6], [7, 8, 9, 10]]
2348        let a = ListViewArray::new(
2349            list_field,
2350            vec![0, 1, 0, 3, 6].into(),
2351            vec![1, 2, 0, 3, 4].into(),
2352            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
2353            None,
2354        );
2355
2356        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2357
2358        assert_eq!(batch.column(0).null_count(), 0);
2359
2360        roundtrip(batch, None);
2361    }
2362
2363    #[test]
2364    fn arrow_writer_list_view_out_of_order() {
2365        let list_field = Arc::new(Field::new_list_field(DataType::Int32, false));
2366        let schema = Schema::new(vec![Field::new(
2367            "a",
2368            DataType::ListView(list_field.clone()),
2369            false,
2370        )]);
2371
2372        // [[1], [2, 3], [], [7, 8, 9, 10], [4, 5, 6]] - out of order offsets
2373        let a = ListViewArray::new(
2374            list_field,
2375            vec![0, 1, 0, 6, 3].into(),
2376            vec![1, 2, 0, 4, 3].into(),
2377            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
2378            None,
2379        );
2380
2381        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2382
2383        roundtrip(batch, None);
2384    }
2385
2386    #[test]
2387    fn arrow_writer_large_list_view() {
2388        let list_field = Arc::new(Field::new_list_field(DataType::Int32, false));
2389        let schema = Schema::new(vec![Field::new(
2390            "a",
2391            DataType::LargeListView(list_field.clone()),
2392            true,
2393        )]);
2394
2395        //  [[1], [2, 3], null, [4, 5, 6], [7, 8, 9, 10]]
2396        let a = LargeListViewArray::new(
2397            list_field,
2398            vec![0i64, 1, 0, 3, 6].into(),
2399            vec![1i64, 2, 0, 3, 4].into(),
2400            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
2401            Some(vec![true, true, false, true, true].into()),
2402        );
2403
2404        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2405
2406        assert_eq!(batch.column(0).null_count(), 1);
2407
2408        roundtrip(batch, None);
2409    }
2410
2411    #[test]
2412    fn arrow_writer_list_view_with_struct() {
2413        // Test ListView containing Struct: ListView<Struct<Int32, Utf8>>
2414        let struct_fields = Fields::from(vec![
2415            Field::new("id", DataType::Int32, false),
2416            Field::new("name", DataType::Utf8, false),
2417        ]);
2418        let struct_type = DataType::Struct(struct_fields.clone());
2419        let list_field = Arc::new(Field::new("item", struct_type.clone(), false));
2420
2421        let schema = Schema::new(vec![Field::new(
2422            "a",
2423            DataType::ListView(list_field.clone()),
2424            true,
2425        )]);
2426
2427        // Create struct values
2428        let id_array = Int32Array::from(vec![1, 2, 3, 4, 5]);
2429        let name_array = StringArray::from(vec!["a", "b", "c", "d", "e"]);
2430        let struct_array = StructArray::new(
2431            struct_fields,
2432            vec![Arc::new(id_array), Arc::new(name_array)],
2433            None,
2434        );
2435
2436        // Create ListView: [{1, "a"}, {2, "b"}], null, [{3, "c"}, {4, "d"}, {5, "e"}]
2437        let list_view = ListViewArray::new(
2438            list_field,
2439            vec![0, 2, 2].into(), // offsets
2440            vec![2, 0, 3].into(), // sizes
2441            Arc::new(struct_array),
2442            Some(vec![true, false, true].into()),
2443        );
2444
2445        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(list_view)]).unwrap();
2446
2447        roundtrip(batch, None);
2448    }
2449
2450    #[test]
2451    fn arrow_writer_binary() {
2452        let string_field = Field::new("a", DataType::Utf8, false);
2453        let binary_field = Field::new("b", DataType::Binary, false);
2454        let schema = Schema::new(vec![string_field, binary_field]);
2455
2456        let raw_string_values = vec!["foo", "bar", "baz", "quux"];
2457        let raw_binary_values = [
2458            b"foo".to_vec(),
2459            b"bar".to_vec(),
2460            b"baz".to_vec(),
2461            b"quux".to_vec(),
2462        ];
2463        let raw_binary_value_refs = raw_binary_values
2464            .iter()
2465            .map(|x| x.as_slice())
2466            .collect::<Vec<_>>();
2467
2468        let string_values = StringArray::from(raw_string_values.clone());
2469        let binary_values = BinaryArray::from(raw_binary_value_refs);
2470        let batch = RecordBatch::try_new(
2471            Arc::new(schema),
2472            vec![Arc::new(string_values), Arc::new(binary_values)],
2473        )
2474        .unwrap();
2475
2476        roundtrip(batch, Some(SMALL_SIZE / 2));
2477    }
2478
2479    #[test]
2480    fn arrow_writer_binary_view() {
2481        let string_field = Field::new("a", DataType::Utf8View, false);
2482        let binary_field = Field::new("b", DataType::BinaryView, false);
2483        let nullable_string_field = Field::new("a", DataType::Utf8View, true);
2484        let schema = Schema::new(vec![string_field, binary_field, nullable_string_field]);
2485
2486        let raw_string_values = vec!["foo", "bar", "large payload over 12 bytes", "lulu"];
2487        let raw_binary_values = vec![
2488            b"foo".to_vec(),
2489            b"bar".to_vec(),
2490            b"large payload over 12 bytes".to_vec(),
2491            b"lulu".to_vec(),
2492        ];
2493        let nullable_string_values =
2494            vec![Some("foo"), None, Some("large payload over 12 bytes"), None];
2495
2496        let string_view_values = StringViewArray::from(raw_string_values);
2497        let binary_view_values = BinaryViewArray::from_iter_values(raw_binary_values);
2498        let nullable_string_view_values = StringViewArray::from(nullable_string_values);
2499        let batch = RecordBatch::try_new(
2500            Arc::new(schema),
2501            vec![
2502                Arc::new(string_view_values),
2503                Arc::new(binary_view_values),
2504                Arc::new(nullable_string_view_values),
2505            ],
2506        )
2507        .unwrap();
2508
2509        roundtrip(batch.clone(), Some(SMALL_SIZE / 2));
2510        roundtrip(batch, None);
2511    }
2512
2513    #[test]
2514    fn arrow_writer_binary_view_long_value() {
2515        let string_field = Field::new("a", DataType::Utf8View, false);
2516        let binary_field = Field::new("b", DataType::BinaryView, false);
2517        let schema = Schema::new(vec![string_field, binary_field]);
2518
2519        // There is special case validation for long values (greater than 128)
2520        // 128 encodes as 0x80 0x00 0x00 0x00 in little endian, which should
2521        // trigger the long-string UTF-8 validation branch in the plain decoder.
2522        let long = "a".repeat(128);
2523        let raw_string_values = vec!["foo", long.as_str(), "bar"];
2524        let raw_binary_values = vec![b"foo".to_vec(), long.as_bytes().to_vec(), b"bar".to_vec()];
2525
2526        let string_view_values: ArrayRef = Arc::new(StringViewArray::from(raw_string_values));
2527        let binary_view_values: ArrayRef =
2528            Arc::new(BinaryViewArray::from_iter_values(raw_binary_values));
2529
2530        RoundTripTest::new(Arc::clone(&string_view_values))
2531            .with_nullable(false)
2532            .run();
2533        RoundTripTest::new(Arc::clone(&binary_view_values))
2534            .with_nullable(false)
2535            .run();
2536
2537        let batch = RecordBatch::try_new(
2538            Arc::new(schema),
2539            vec![string_view_values, binary_view_values],
2540        )
2541        .unwrap();
2542
2543        // Disable dictionary to exercise plain encoding paths in the reader.
2544        for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] {
2545            let props = WriterProperties::builder()
2546                .set_writer_version(version)
2547                .set_dictionary_enabled(false)
2548                .build();
2549            roundtrip_opts(&batch, props);
2550        }
2551    }
2552
2553    fn get_decimal_batch(precision: u8, scale: i8) -> RecordBatch {
2554        let decimal_field = Field::new("a", DataType::Decimal128(precision, scale), false);
2555        let schema = Schema::new(vec![decimal_field]);
2556
2557        let decimal_values = vec![10_000, 50_000, 0, -100]
2558            .into_iter()
2559            .map(Some)
2560            .collect::<Decimal128Array>()
2561            .with_precision_and_scale(precision, scale)
2562            .unwrap();
2563
2564        RecordBatch::try_new(Arc::new(schema), vec![Arc::new(decimal_values)]).unwrap()
2565    }
2566
2567    #[test]
2568    fn arrow_writer_decimal() {
2569        // int32 to store the decimal value
2570        let batch_int32_decimal = get_decimal_batch(5, 2);
2571        roundtrip(batch_int32_decimal, Some(SMALL_SIZE / 2));
2572        // int64 to store the decimal value
2573        let batch_int64_decimal = get_decimal_batch(12, 2);
2574        roundtrip(batch_int64_decimal, Some(SMALL_SIZE / 2));
2575        // fixed_length_byte_array to store the decimal value
2576        let batch_fixed_len_byte_array_decimal = get_decimal_batch(30, 2);
2577        roundtrip(batch_fixed_len_byte_array_decimal, Some(SMALL_SIZE / 2));
2578    }
2579
2580    #[test]
2581    fn arrow_writer_complex() {
2582        // define schema
2583        let struct_field_d = Arc::new(Field::new("d", DataType::Float64, true));
2584        let struct_field_f = Arc::new(Field::new("f", DataType::Float32, true));
2585        let struct_field_g = Arc::new(Field::new_list(
2586            "g",
2587            Field::new_list_field(DataType::Int16, true),
2588            false,
2589        ));
2590        let struct_field_h = Arc::new(Field::new_list(
2591            "h",
2592            Field::new_list_field(DataType::Int16, false),
2593            true,
2594        ));
2595        let struct_field_e = Arc::new(Field::new_struct(
2596            "e",
2597            vec![
2598                struct_field_f.clone(),
2599                struct_field_g.clone(),
2600                struct_field_h.clone(),
2601            ],
2602            false,
2603        ));
2604        let schema = Schema::new(vec![
2605            Field::new("a", DataType::Int32, false),
2606            Field::new("b", DataType::Int32, true),
2607            Field::new_struct(
2608                "c",
2609                vec![struct_field_d.clone(), struct_field_e.clone()],
2610                false,
2611            ),
2612        ]);
2613
2614        // create some data
2615        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2616        let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2617        let d = Float64Array::from(vec![None, None, None, Some(1.0), None]);
2618        let f = Float32Array::from(vec![Some(0.0), None, Some(333.3), None, Some(5.25)]);
2619
2620        let g_value = Int16Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
2621
2622        // Construct a buffer for value offsets, for the nested array:
2623        //  [[1], [2, 3], [], [4, 5, 6], [7, 8, 9, 10]]
2624        let g_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice());
2625
2626        // Construct a list array from the above two
2627        let g_list_data = ArrayData::builder(struct_field_g.data_type().clone())
2628            .len(5)
2629            .add_buffer(g_value_offsets.clone())
2630            .add_child_data(g_value.to_data())
2631            .build()
2632            .unwrap();
2633        let g = ListArray::from(g_list_data);
2634        // The difference between g and h is that h has a null bitmap
2635        let h_list_data = ArrayData::builder(struct_field_h.data_type().clone())
2636            .len(5)
2637            .add_buffer(g_value_offsets)
2638            .add_child_data(g_value.to_data())
2639            .null_bit_buffer(Some(Buffer::from([0b00011011])))
2640            .build()
2641            .unwrap();
2642        let h = ListArray::from(h_list_data);
2643
2644        let e = StructArray::from(vec![
2645            (struct_field_f, Arc::new(f) as ArrayRef),
2646            (struct_field_g, Arc::new(g) as ArrayRef),
2647            (struct_field_h, Arc::new(h) as ArrayRef),
2648        ]);
2649
2650        let c = StructArray::from(vec![
2651            (struct_field_d, Arc::new(d) as ArrayRef),
2652            (struct_field_e, Arc::new(e) as ArrayRef),
2653        ]);
2654
2655        // build a record batch
2656        let batch = RecordBatch::try_new(
2657            Arc::new(schema),
2658            vec![Arc::new(a), Arc::new(b), Arc::new(c)],
2659        )
2660        .unwrap();
2661
2662        roundtrip(batch.clone(), Some(SMALL_SIZE / 2));
2663        roundtrip(batch, Some(SMALL_SIZE / 3));
2664    }
2665
2666    #[test]
2667    fn arrow_writer_complex_mixed() {
2668        // This test was added while investigating https://github.com/apache/arrow-rs/issues/244.
2669        // It was subsequently fixed while investigating https://github.com/apache/arrow-rs/issues/245.
2670
2671        // define schema
2672        let offset_field = Arc::new(Field::new("offset", DataType::Int32, false));
2673        let partition_field = Arc::new(Field::new("partition", DataType::Int64, true));
2674        let topic_field = Arc::new(Field::new("topic", DataType::Utf8, true));
2675        let schema = Schema::new(vec![Field::new(
2676            "some_nested_object",
2677            DataType::Struct(Fields::from(vec![
2678                offset_field.clone(),
2679                partition_field.clone(),
2680                topic_field.clone(),
2681            ])),
2682            false,
2683        )]);
2684
2685        // create some data
2686        let offset = Int32Array::from(vec![1, 2, 3, 4, 5]);
2687        let partition = Int64Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2688        let topic = StringArray::from(vec![Some("A"), None, Some("A"), Some(""), None]);
2689
2690        let some_nested_object = StructArray::from(vec![
2691            (offset_field, Arc::new(offset) as ArrayRef),
2692            (partition_field, Arc::new(partition) as ArrayRef),
2693            (topic_field, Arc::new(topic) as ArrayRef),
2694        ]);
2695
2696        // build a record batch
2697        let batch =
2698            RecordBatch::try_new(Arc::new(schema), vec![Arc::new(some_nested_object)]).unwrap();
2699
2700        roundtrip(batch, Some(SMALL_SIZE / 2));
2701    }
2702
2703    #[test]
2704    fn arrow_writer_map() {
2705        // Note: we are using the JSON Arrow reader for brevity
2706        let json_content = r#"
2707        {"stocks":{"long": "$AAA", "short": "$BBB"}}
2708        {"stocks":{"long": null, "long": "$CCC", "short": null}}
2709        {"stocks":{"hedged": "$YYY", "long": null, "short": "$D"}}
2710        "#;
2711        let entries_struct_type = DataType::Struct(Fields::from(vec![
2712            Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
2713            Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
2714        ]));
2715        let stocks_field = Field::new(
2716            "stocks",
2717            DataType::Map(
2718                Arc::new(Field::new(
2719                    Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
2720                    entries_struct_type,
2721                    false,
2722                )),
2723                false,
2724            ),
2725            true,
2726        );
2727        let schema = Arc::new(Schema::new(vec![stocks_field]));
2728        let builder = arrow::json::ReaderBuilder::new(schema).with_batch_size(64);
2729        let mut reader = builder.build(std::io::Cursor::new(json_content)).unwrap();
2730
2731        let batch = reader.next().unwrap().unwrap();
2732        roundtrip(batch, None);
2733    }
2734
2735    #[test]
2736    fn arrow_writer_2_level_struct() {
2737        // tests writing <struct<struct<primitive>>
2738        let field_c = Field::new("c", DataType::Int32, true);
2739        let field_b = Field::new("b", DataType::Struct(vec![field_c].into()), true);
2740        let type_a = DataType::Struct(vec![field_b.clone()].into());
2741        let field_a = Field::new("a", type_a, true);
2742        let schema = Schema::new(vec![field_a.clone()]);
2743
2744        // create data
2745        let c = Int32Array::from(vec![Some(1), None, Some(3), None, None, Some(6)]);
2746        let b_data = ArrayDataBuilder::new(field_b.data_type().clone())
2747            .len(6)
2748            .null_bit_buffer(Some(Buffer::from([0b00100111])))
2749            .add_child_data(c.into_data())
2750            .build()
2751            .unwrap();
2752        let b = StructArray::from(b_data);
2753        let a_data = ArrayDataBuilder::new(field_a.data_type().clone())
2754            .len(6)
2755            .null_bit_buffer(Some(Buffer::from([0b00101111])))
2756            .add_child_data(b.into_data())
2757            .build()
2758            .unwrap();
2759        let a = StructArray::from(a_data);
2760
2761        assert_eq!(a.null_count(), 1);
2762        assert_eq!(a.column(0).null_count(), 2);
2763
2764        // build a racord batch
2765        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2766
2767        roundtrip(batch, Some(SMALL_SIZE / 2));
2768    }
2769
2770    #[test]
2771    fn arrow_writer_2_level_struct_non_null() {
2772        // tests writing <struct<struct<primitive>>
2773        let field_c = Field::new("c", DataType::Int32, false);
2774        let type_b = DataType::Struct(vec![field_c].into());
2775        let field_b = Field::new("b", type_b.clone(), false);
2776        let type_a = DataType::Struct(vec![field_b].into());
2777        let field_a = Field::new("a", type_a.clone(), false);
2778        let schema = Schema::new(vec![field_a]);
2779
2780        // create data
2781        let c = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
2782        let b_data = ArrayDataBuilder::new(type_b)
2783            .len(6)
2784            .add_child_data(c.into_data())
2785            .build()
2786            .unwrap();
2787        let b = StructArray::from(b_data);
2788        let a_data = ArrayDataBuilder::new(type_a)
2789            .len(6)
2790            .add_child_data(b.into_data())
2791            .build()
2792            .unwrap();
2793        let a = StructArray::from(a_data);
2794
2795        assert_eq!(a.null_count(), 0);
2796        assert_eq!(a.column(0).null_count(), 0);
2797
2798        // build a racord batch
2799        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2800
2801        roundtrip(batch, Some(SMALL_SIZE / 2));
2802    }
2803
2804    #[test]
2805    fn arrow_writer_2_level_struct_mixed_null() {
2806        // tests writing <struct<struct<primitive>>
2807        let field_c = Field::new("c", DataType::Int32, false);
2808        let type_b = DataType::Struct(vec![field_c].into());
2809        let field_b = Field::new("b", type_b.clone(), true);
2810        let type_a = DataType::Struct(vec![field_b].into());
2811        let field_a = Field::new("a", type_a.clone(), false);
2812        let schema = Schema::new(vec![field_a]);
2813
2814        // create data
2815        let c = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
2816        let b_data = ArrayDataBuilder::new(type_b)
2817            .len(6)
2818            .null_bit_buffer(Some(Buffer::from([0b00100111])))
2819            .add_child_data(c.into_data())
2820            .build()
2821            .unwrap();
2822        let b = StructArray::from(b_data);
2823        // a intentionally has no null buffer, to test that this is handled correctly
2824        let a_data = ArrayDataBuilder::new(type_a)
2825            .len(6)
2826            .add_child_data(b.into_data())
2827            .build()
2828            .unwrap();
2829        let a = StructArray::from(a_data);
2830
2831        assert_eq!(a.null_count(), 0);
2832        assert_eq!(a.column(0).null_count(), 2);
2833
2834        // build a racord batch
2835        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2836
2837        roundtrip(batch, Some(SMALL_SIZE / 2));
2838    }
2839
2840    #[test]
2841    fn arrow_writer_2_level_struct_mixed_null_2() {
2842        // tests writing <struct<struct<primitive>>, where the primitive columns are non-null.
2843        let field_c = Field::new("c", DataType::Int32, false);
2844        let field_d = Field::new("d", DataType::FixedSizeBinary(4), false);
2845        let field_e = Field::new(
2846            "e",
2847            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
2848            false,
2849        );
2850
2851        let field_b = Field::new(
2852            "b",
2853            DataType::Struct(vec![field_c, field_d, field_e].into()),
2854            false,
2855        );
2856        let type_a = DataType::Struct(vec![field_b.clone()].into());
2857        let field_a = Field::new("a", type_a, true);
2858        let schema = Schema::new(vec![field_a.clone()]);
2859
2860        // create data
2861        let c = Int32Array::from_iter_values(0..6);
2862        let d = FixedSizeBinaryArray::try_from_iter(
2863            ["aaaa", "bbbb", "cccc", "dddd", "eeee", "ffff"].into_iter(),
2864        )
2865        .expect("four byte values");
2866        let e = Int32DictionaryArray::from_iter(["one", "two", "three", "four", "five", "one"]);
2867        let b_data = ArrayDataBuilder::new(field_b.data_type().clone())
2868            .len(6)
2869            .add_child_data(c.into_data())
2870            .add_child_data(d.into_data())
2871            .add_child_data(e.into_data())
2872            .build()
2873            .unwrap();
2874        let b = StructArray::from(b_data);
2875        let a_data = ArrayDataBuilder::new(field_a.data_type().clone())
2876            .len(6)
2877            .null_bit_buffer(Some(Buffer::from([0b00100101])))
2878            .add_child_data(b.into_data())
2879            .build()
2880            .unwrap();
2881        let a = StructArray::from(a_data);
2882
2883        assert_eq!(a.null_count(), 3);
2884        assert_eq!(a.column(0).null_count(), 0);
2885
2886        // build a record batch
2887        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2888
2889        roundtrip(batch, Some(SMALL_SIZE / 2));
2890    }
2891
2892    #[test]
2893    fn test_fixed_size_binary_in_dict() {
2894        fn test_fixed_size_binary_in_dict_inner<K>()
2895        where
2896            K: ArrowDictionaryKeyType,
2897            K::Native: FromPrimitive + ToPrimitive + TryFrom<u8>,
2898            <<K as arrow_array::ArrowPrimitiveType>::Native as TryFrom<u8>>::Error: std::fmt::Debug,
2899        {
2900            let field = Field::new(
2901                "a",
2902                DataType::Dictionary(
2903                    Box::new(K::DATA_TYPE),
2904                    Box::new(DataType::FixedSizeBinary(4)),
2905                ),
2906                false,
2907            );
2908            let schema = Schema::new(vec![field]);
2909
2910            let keys: Vec<K::Native> = vec![
2911                K::Native::try_from(0u8).unwrap(),
2912                K::Native::try_from(0u8).unwrap(),
2913                K::Native::try_from(1u8).unwrap(),
2914            ];
2915            let keys = PrimitiveArray::<K>::from_iter_values(keys);
2916            let values = FixedSizeBinaryArray::try_from_iter(
2917                vec![vec![0, 0, 0, 0], vec![1, 1, 1, 1]].into_iter(),
2918            )
2919            .unwrap();
2920
2921            let data = DictionaryArray::<K>::new(keys, Arc::new(values));
2922            let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(data)]).unwrap();
2923            roundtrip(batch, None);
2924        }
2925
2926        test_fixed_size_binary_in_dict_inner::<UInt8Type>();
2927        test_fixed_size_binary_in_dict_inner::<UInt16Type>();
2928        test_fixed_size_binary_in_dict_inner::<UInt32Type>();
2929        test_fixed_size_binary_in_dict_inner::<UInt16Type>();
2930        test_fixed_size_binary_in_dict_inner::<Int8Type>();
2931        test_fixed_size_binary_in_dict_inner::<Int16Type>();
2932        test_fixed_size_binary_in_dict_inner::<Int32Type>();
2933        test_fixed_size_binary_in_dict_inner::<Int64Type>();
2934    }
2935
2936    #[test]
2937    fn test_empty_dict() {
2938        let struct_fields = Fields::from(vec![Field::new(
2939            "dict",
2940            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
2941            false,
2942        )]);
2943
2944        let schema = Schema::new(vec![Field::new_struct(
2945            "struct",
2946            struct_fields.clone(),
2947            true,
2948        )]);
2949        let dictionary = Arc::new(DictionaryArray::new(
2950            Int32Array::new_null(5),
2951            Arc::new(StringArray::new_null(0)),
2952        ));
2953
2954        let s = StructArray::new(
2955            struct_fields,
2956            vec![dictionary],
2957            Some(NullBuffer::new_null(5)),
2958        );
2959
2960        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(s)]).unwrap();
2961        roundtrip(batch, None);
2962    }
2963    #[test]
2964    fn arrow_writer_page_size() {
2965        let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)]));
2966
2967        let mut builder = StringBuilder::with_capacity(100, 329 * 10_000);
2968
2969        // Generate an array of 10 unique 10 character string
2970        for i in 0..10 {
2971            let value = i
2972                .to_string()
2973                .repeat(10)
2974                .chars()
2975                .take(10)
2976                .collect::<String>();
2977
2978            builder.append_value(value);
2979        }
2980
2981        let array = Arc::new(builder.finish());
2982
2983        let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
2984
2985        let file = tempfile::tempfile().unwrap();
2986
2987        // Set everything very low so we fallback to PLAIN encoding after the first row
2988        let props = WriterProperties::builder()
2989            .set_data_page_size_limit(1)
2990            .set_dictionary_page_size_limit(1)
2991            .set_write_batch_size(1)
2992            .build();
2993
2994        let mut writer =
2995            ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props))
2996                .expect("Unable to write file");
2997        writer.write(&batch).unwrap();
2998        writer.close().unwrap();
2999
3000        let options = ReadOptionsBuilder::new().with_page_index().build();
3001        let reader =
3002            SerializedFileReader::new_with_options(file.try_clone().unwrap(), options).unwrap();
3003
3004        let column = reader.metadata().row_group(0).columns();
3005
3006        assert_eq!(column.len(), 1);
3007
3008        // We should write one row before falling back to PLAIN encoding so there should still be a
3009        // dictionary page.
3010        assert!(
3011            column[0].dictionary_page_offset().is_some(),
3012            "Expected a dictionary page"
3013        );
3014
3015        assert!(reader.metadata().offset_index().is_some());
3016        let offset_indexes = &reader.metadata().offset_index().unwrap()[0];
3017
3018        let page_locations = offset_indexes[0].page_locations.clone();
3019
3020        // We should fallback to PLAIN encoding after the first row and our max page size is 1 bytes
3021        // so we expect one dictionary encoded page and then a page per row thereafter.
3022        assert_eq!(
3023            page_locations.len(),
3024            10,
3025            "Expected 10 pages but got {page_locations:#?}"
3026        );
3027    }
3028
3029    #[test]
3030    fn arrow_writer_float_nans() {
3031        let f16_field = Field::new("a", DataType::Float16, false);
3032        let f32_field = Field::new("b", DataType::Float32, false);
3033        let f64_field = Field::new("c", DataType::Float64, false);
3034        let schema = Schema::new(vec![f16_field, f32_field, f64_field]);
3035
3036        let f16_values = (0..MEDIUM_SIZE)
3037            .map(|i| {
3038                Some(if i % 2 == 0 {
3039                    f16::NAN
3040                } else {
3041                    f16::from_f32(i as f32)
3042                })
3043            })
3044            .collect::<Float16Array>();
3045
3046        let f32_values = (0..MEDIUM_SIZE)
3047            .map(|i| Some(if i % 2 == 0 { f32::NAN } else { i as f32 }))
3048            .collect::<Float32Array>();
3049
3050        let f64_values = (0..MEDIUM_SIZE)
3051            .map(|i| Some(if i % 2 == 0 { f64::NAN } else { i as f64 }))
3052            .collect::<Float64Array>();
3053
3054        let batch = RecordBatch::try_new(
3055            Arc::new(schema),
3056            vec![
3057                Arc::new(f16_values),
3058                Arc::new(f32_values),
3059                Arc::new(f64_values),
3060            ],
3061        )
3062        .unwrap();
3063
3064        roundtrip(batch, None);
3065    }
3066
3067    const SMALL_SIZE: usize = 7;
3068    const MEDIUM_SIZE: usize = 63;
3069
3070    // Write the batch to parquet and read it back out, ensuring
3071    // that what comes out is the same as what was written in
3072    fn roundtrip(expected_batch: RecordBatch, max_row_group_size: Option<usize>) -> Vec<Bytes> {
3073        let mut files = vec![];
3074        for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] {
3075            let mut props = WriterProperties::builder().set_writer_version(version);
3076
3077            if let Some(size) = max_row_group_size {
3078                props = props.set_max_row_group_row_count(Some(size))
3079            }
3080
3081            let props = props.build();
3082            files.push(roundtrip_opts(&expected_batch, props))
3083        }
3084        files
3085    }
3086
3087    // Round trip the specified record batch with the specified writer properties,
3088    // to an in-memory file, and validate the arrays using the specified function.
3089    // Returns the in-memory file.
3090    fn roundtrip_opts_with_array_validation<F>(
3091        expected_batch: &RecordBatch,
3092        props: WriterProperties,
3093        validate: F,
3094    ) -> Bytes
3095    where
3096        F: Fn(&ArrayData, &ArrayData),
3097    {
3098        let mut file = vec![];
3099
3100        let mut writer = ArrowWriter::try_new(&mut file, expected_batch.schema(), Some(props))
3101            .expect("Unable to write file");
3102        writer.write(expected_batch).unwrap();
3103        writer.close().unwrap();
3104
3105        let file = Bytes::from(file);
3106        let mut record_batch_reader =
3107            ParquetRecordBatchReader::try_new(file.clone(), 1024).unwrap();
3108
3109        let actual_batch = record_batch_reader
3110            .next()
3111            .expect("No batch found")
3112            .expect("Unable to get batch");
3113
3114        assert_eq!(expected_batch.schema(), actual_batch.schema());
3115        assert_eq!(expected_batch.num_columns(), actual_batch.num_columns());
3116        assert_eq!(expected_batch.num_rows(), actual_batch.num_rows());
3117        for i in 0..expected_batch.num_columns() {
3118            let expected_data = expected_batch.column(i).to_data();
3119            let actual_data = actual_batch.column(i).to_data();
3120            validate(&expected_data, &actual_data);
3121        }
3122
3123        file
3124    }
3125
3126    fn roundtrip_opts(expected_batch: &RecordBatch, props: WriterProperties) -> Bytes {
3127        roundtrip_opts_with_array_validation(expected_batch, props, |a, b| {
3128            a.validate_full().expect("valid expected data");
3129            b.validate_full().expect("valid actual data");
3130            assert_eq!(a, b)
3131        })
3132    }
3133
3134    /// Round trip testing fixture:
3135    ///
3136    /// Tests based on this fixture write data to parquet and then read it back.
3137    struct RoundTripTest {
3138        values: ArrayRef,
3139        /// Optionally supplied schema
3140        schema: Option<SchemaRef>,
3141        /// If the created schema should be nullable. Defaults to true. Ignored
3142        /// if schema is set to Some.
3143        nullable: bool,
3144        bloom_filter: bool,
3145        bloom_filter_ndv: Option<u64>,
3146        bloom_filter_position: BloomFilterPosition,
3147    }
3148
3149    impl RoundTripTest {
3150        /// Create a test for round tripping values with a nullable schema
3151        fn new(values: ArrayRef) -> Self {
3152            Self {
3153                values,
3154                schema: None,
3155                nullable: true,
3156                bloom_filter: false,
3157                bloom_filter_ndv: None,
3158                bloom_filter_position: BloomFilterPosition::AfterRowGroup,
3159            }
3160        }
3161
3162        /// Set the schema
3163        fn with_schema(mut self, schema: SchemaRef) -> Self {
3164            self.schema = Some(schema);
3165            self
3166        }
3167
3168        /// Set the nullable flag
3169        fn with_nullable(mut self, nullable: bool) -> Self {
3170            self.nullable = nullable;
3171            self
3172        }
3173
3174        /// Set bloom filter
3175        fn with_bloom_filter(mut self, bloom_filter: bool) -> Self {
3176            self.bloom_filter = bloom_filter;
3177            self
3178        }
3179
3180        /// Set bloom filter max ndv
3181        fn with_bloom_filter_ndv(mut self, bloom_filter_ndv: u64) -> Self {
3182            self.bloom_filter_ndv = Some(bloom_filter_ndv);
3183            self
3184        }
3185
3186        /// Set bloom filter position
3187        fn with_bloom_filter_position(
3188            mut self,
3189            bloom_filter_position: BloomFilterPosition,
3190        ) -> Self {
3191            self.bloom_filter_position = bloom_filter_position;
3192            self
3193        }
3194
3195        /// Run the test specified by the options, returning the encoded Parquet bytes
3196        fn run(self) -> Vec<Bytes> {
3197            let RoundTripTest {
3198                values,
3199                schema,
3200                nullable,
3201                bloom_filter,
3202                bloom_filter_ndv,
3203                bloom_filter_position,
3204            } = self;
3205
3206            let schema = schema.unwrap_or_else(|| {
3207                let data_type = values.data_type().clone();
3208                Arc::new(Schema::new(vec![Field::new("col", data_type, nullable)]))
3209            });
3210
3211            let encodings = match values.data_type() {
3212                DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary | DataType::LargeBinary => {
3213                    vec![
3214                        Encoding::PLAIN,
3215                        Encoding::DELTA_BYTE_ARRAY,
3216                        Encoding::DELTA_LENGTH_BYTE_ARRAY,
3217                    ]
3218                }
3219                DataType::Int64
3220                | DataType::Int32
3221                | DataType::Int16
3222                | DataType::Int8
3223                | DataType::UInt64
3224                | DataType::UInt32
3225                | DataType::UInt16
3226                | DataType::UInt8 => vec![
3227                    Encoding::PLAIN,
3228                    Encoding::DELTA_BINARY_PACKED,
3229                    Encoding::BYTE_STREAM_SPLIT,
3230                ],
3231                DataType::Float32 | DataType::Float64 => {
3232                    vec![Encoding::PLAIN, Encoding::BYTE_STREAM_SPLIT]
3233                }
3234                _ => vec![Encoding::PLAIN],
3235            };
3236
3237            let expected_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
3238
3239            let row_group_sizes = [1024, SMALL_SIZE, SMALL_SIZE / 2, SMALL_SIZE / 2 + 1, 10];
3240
3241            let mut files = vec![];
3242            for dictionary_size in [0, 1, 1024] {
3243                for encoding in &encodings {
3244                    for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] {
3245                        for row_group_size in row_group_sizes {
3246                            let mut builder = WriterProperties::builder()
3247                                .set_writer_version(version)
3248                                .set_max_row_group_row_count(Some(row_group_size))
3249                                .set_dictionary_enabled(dictionary_size != 0)
3250                                .set_dictionary_page_size_limit(dictionary_size.max(1))
3251                                .set_encoding(*encoding)
3252                                .set_bloom_filter_enabled(bloom_filter)
3253                                .set_bloom_filter_position(bloom_filter_position);
3254                            if let Some(ndv) = bloom_filter_ndv {
3255                                builder = builder.set_bloom_filter_max_ndv(ndv);
3256                            }
3257                            let props = builder.build();
3258
3259                            files.push(roundtrip_opts(&expected_batch, props))
3260                        }
3261                    }
3262                }
3263            }
3264            files
3265        }
3266    }
3267
3268    fn values_required<A, I>(iter: I) -> Vec<Bytes>
3269    where
3270        A: From<Vec<I::Item>> + Array + 'static,
3271        I: IntoIterator,
3272    {
3273        let raw_values: Vec<_> = iter.into_iter().collect();
3274        let values = Arc::new(A::from(raw_values));
3275        RoundTripTest::new(values).with_nullable(false).run()
3276    }
3277
3278    fn values_optional<A, I>(iter: I) -> Vec<Bytes>
3279    where
3280        A: From<Vec<Option<I::Item>>> + Array + 'static,
3281        I: IntoIterator,
3282    {
3283        let optional_raw_values: Vec<_> = iter
3284            .into_iter()
3285            .enumerate()
3286            .map(|(i, v)| if i % 2 == 0 { None } else { Some(v) })
3287            .collect();
3288        let optional_values = Arc::new(A::from(optional_raw_values));
3289        RoundTripTest::new(optional_values).run()
3290    }
3291
3292    fn required_and_optional<A, I>(iter: I)
3293    where
3294        A: From<Vec<I::Item>> + From<Vec<Option<I::Item>>> + Array + 'static,
3295        I: IntoIterator + Clone,
3296    {
3297        values_required::<A, I>(iter.clone());
3298        values_optional::<A, I>(iter);
3299    }
3300
3301    fn check_bloom_filter<T: AsBytes>(
3302        files: Vec<Bytes>,
3303        file_column: String,
3304        positive_values: Vec<T>,
3305        negative_values: Vec<T>,
3306    ) {
3307        files.into_iter().take(1).for_each(|file| {
3308            let file_reader = SerializedFileReader::new_with_options(
3309                file,
3310                ReadOptionsBuilder::new()
3311                    .with_reader_properties(
3312                        ReaderProperties::builder()
3313                            .set_read_bloom_filter(true)
3314                            .build(),
3315                    )
3316                    .build(),
3317            )
3318            .expect("Unable to open file as Parquet");
3319            let metadata = file_reader.metadata();
3320
3321            // Gets bloom filters from all row groups.
3322            let mut bloom_filters: Vec<_> = vec![];
3323            for (ri, row_group) in metadata.row_groups().iter().enumerate() {
3324                if let Some((column_index, _)) = row_group
3325                    .columns()
3326                    .iter()
3327                    .enumerate()
3328                    .find(|(_, column)| column.column_path().string() == file_column)
3329                {
3330                    let row_group_reader = file_reader
3331                        .get_row_group(ri)
3332                        .expect("Unable to read row group");
3333                    if let Some(sbbf) = row_group_reader.get_column_bloom_filter(column_index) {
3334                        bloom_filters.push(sbbf.clone());
3335                    } else {
3336                        panic!("No bloom filter for column named {file_column} found");
3337                    }
3338                } else {
3339                    panic!("No column named {file_column} found");
3340                }
3341            }
3342
3343            positive_values.iter().for_each(|value| {
3344                let found = bloom_filters.iter().find(|sbbf| sbbf.check(value));
3345                assert!(
3346                    found.is_some(),
3347                    "{}",
3348                    format!("Value {:?} should be in bloom filter", value.as_bytes())
3349                );
3350            });
3351
3352            negative_values.iter().for_each(|value| {
3353                let found = bloom_filters.iter().find(|sbbf| sbbf.check(value));
3354                assert!(
3355                    found.is_none(),
3356                    "{}",
3357                    format!("Value {:?} should not be in bloom filter", value.as_bytes())
3358                );
3359            });
3360        });
3361    }
3362
3363    #[test]
3364    fn all_null_primitive_single_column() {
3365        let values = Arc::new(Int32Array::from(vec![None; SMALL_SIZE]));
3366        RoundTripTest::new(values).run();
3367    }
3368    #[test]
3369    fn null_single_column() {
3370        let values = Arc::new(NullArray::new(SMALL_SIZE));
3371        RoundTripTest::new(values).run();
3372        // null arrays are always nullable, a test with non-nullable nulls fails
3373    }
3374
3375    #[test]
3376    fn bool_single_column() {
3377        required_and_optional::<BooleanArray, _>(
3378            [true, false].iter().cycle().copied().take(SMALL_SIZE),
3379        );
3380    }
3381
3382    #[test]
3383    fn bool_large_single_column() {
3384        let values = Arc::new(
3385            [None, Some(true), Some(false)]
3386                .iter()
3387                .cycle()
3388                .copied()
3389                .take(200_000)
3390                .collect::<BooleanArray>(),
3391        );
3392        let schema = Schema::new(vec![Field::new("col", values.data_type().clone(), true)]);
3393        let expected_batch = RecordBatch::try_new(Arc::new(schema), vec![values]).unwrap();
3394        let file = tempfile::tempfile().unwrap();
3395
3396        let mut writer =
3397            ArrowWriter::try_new(file.try_clone().unwrap(), expected_batch.schema(), None)
3398                .expect("Unable to write file");
3399        writer.write(&expected_batch).unwrap();
3400        writer.close().unwrap();
3401    }
3402
3403    #[test]
3404    fn check_page_offset_index_with_nan() {
3405        let values = Arc::new(Float64Array::from(vec![f64::NAN; 10]));
3406        let schema = Schema::new(vec![Field::new("col", DataType::Float64, true)]);
3407        let batch = RecordBatch::try_new(Arc::new(schema), vec![values]).unwrap();
3408
3409        let mut out = Vec::with_capacity(1024);
3410        let mut writer =
3411            ArrowWriter::try_new(&mut out, batch.schema(), None).expect("Unable to write file");
3412        writer.write(&batch).unwrap();
3413        let file_meta_data = writer.close().unwrap();
3414        for row_group in file_meta_data.row_groups() {
3415            for column in row_group.columns() {
3416                assert!(column.offset_index_offset().is_some());
3417                assert!(column.offset_index_length().is_some());
3418                assert!(column.column_index_offset().is_some());
3419                assert!(column.column_index_length().is_some());
3420            }
3421        }
3422        assert!(file_meta_data.column_index().is_some());
3423        if let Some(col_indexes) = file_meta_data.column_index() {
3424            for rg_idx in col_indexes {
3425                for idx in rg_idx {
3426                    assert!(idx.nan_counts().is_some());
3427                    let float_idx = match idx {
3428                        ColumnIndexMetaData::DOUBLE(idx) => idx,
3429                        _ => panic!("expected double statistics"),
3430                    };
3431                    for i in 0..idx.num_pages() as usize {
3432                        assert_eq!(float_idx.nan_count(i), Some(10));
3433                        assert_eq!(
3434                            f64::NAN.total_cmp(float_idx.min_value(i).unwrap()),
3435                            Ordering::Equal
3436                        );
3437                        assert_eq!(
3438                            f64::NAN.total_cmp(float_idx.max_value(i).unwrap()),
3439                            Ordering::Equal
3440                        );
3441                    }
3442                }
3443            }
3444        }
3445    }
3446
3447    #[test]
3448    fn check_page_offset_index_with_mixed_nan() {
3449        let schema = Arc::new(Schema::new(vec![Field::new(
3450            "col",
3451            DataType::Float64,
3452            true,
3453        )]));
3454
3455        let mut out = Vec::with_capacity(1024);
3456        let props = WriterProperties::builder()
3457            .set_data_page_row_count_limit(10)
3458            .build();
3459        let mut writer = ArrowWriter::try_new(&mut out, schema.clone(), Some(props))
3460            .expect("Unable to write file");
3461
3462        // write a page of all NaN (since batch min and max are NaN, global min/max are NaN)
3463        let values = Arc::new(Float64Array::from(vec![f64::NAN; 10]));
3464        let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
3465        writer.write(&batch).unwrap();
3466
3467        // write a page of all -NaN (batch min/max is -NaN, should update global min to -NaN)
3468        let values = Arc::new(Float64Array::from(vec![-f64::NAN; 10]));
3469        let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
3470        writer.write(&batch).unwrap();
3471
3472        // write a page of all 0 (non-NaN should override global min/max, now 0/0)
3473        let values = Arc::new(Float64Array::from(vec![0_f64; 10]));
3474        let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
3475        writer.write(&batch).unwrap();
3476
3477        // write a mixed page (should now have min -1, max 1)
3478        let values = Arc::new(Float64Array::from(vec![
3479            -1.0,
3480            0.0,
3481            f64::NAN,
3482            -f64::NAN,
3483            1.0,
3484        ]));
3485        let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
3486        writer.write(&batch).unwrap();
3487
3488        let file_meta_data = writer.close().unwrap();
3489
3490        // check the column chunk stats are correct
3491        let col_stats = file_meta_data
3492            .row_group(0)
3493            .column(0)
3494            .statistics()
3495            .expect("missing column chunk statistics");
3496
3497        assert_eq!(col_stats.nan_count_opt(), Some(22));
3498        assert_eq!(col_stats.min_bytes_opt(), Some((-1.0f64).as_bytes()));
3499        assert_eq!(col_stats.max_bytes_opt(), Some(1.0f64.as_bytes()));
3500
3501        assert!(file_meta_data.column_index().is_some());
3502        let col_idx = &file_meta_data.column_index().as_ref().unwrap()[0][0];
3503        assert_eq!(col_idx.num_pages(), 4);
3504
3505        // test each page
3506        let float_idx = match col_idx {
3507            ColumnIndexMetaData::DOUBLE(idx) => idx,
3508            _ => panic!("expected double statistics"),
3509        };
3510
3511        assert_eq!(float_idx.nan_counts, Some(vec![10, 10, 0, 2]));
3512        assert_eq!(
3513            f64::NAN.total_cmp(float_idx.min_value(0).unwrap()),
3514            Ordering::Equal
3515        );
3516        assert_eq!(
3517            f64::NAN.total_cmp(float_idx.max_value(0).unwrap()),
3518            Ordering::Equal
3519        );
3520        assert_eq!(
3521            (-f64::NAN).total_cmp(float_idx.min_value(1).unwrap()),
3522            Ordering::Equal
3523        );
3524        assert_eq!(
3525            (-f64::NAN).total_cmp(float_idx.max_value(1).unwrap()),
3526            Ordering::Equal
3527        );
3528        assert_eq!(float_idx.min_value(2), Some(&0.0));
3529        assert_eq!(float_idx.max_value(2), Some(&0.0));
3530        assert_eq!(float_idx.min_value(3), Some(&-1.0));
3531        assert_eq!(float_idx.max_value(3), Some(&1.0));
3532    }
3533
3534    #[test]
3535    fn i8_single_column() {
3536        required_and_optional::<Int8Array, _>(0..SMALL_SIZE as i8);
3537    }
3538
3539    #[test]
3540    fn i16_single_column() {
3541        required_and_optional::<Int16Array, _>(0..SMALL_SIZE as i16);
3542    }
3543
3544    #[test]
3545    fn i32_single_column() {
3546        required_and_optional::<Int32Array, _>(0..SMALL_SIZE as i32);
3547    }
3548
3549    #[test]
3550    fn i64_single_column() {
3551        required_and_optional::<Int64Array, _>(0..SMALL_SIZE as i64);
3552    }
3553
3554    #[test]
3555    fn u8_single_column() {
3556        required_and_optional::<UInt8Array, _>(0..SMALL_SIZE as u8);
3557    }
3558
3559    #[test]
3560    fn u16_single_column() {
3561        required_and_optional::<UInt16Array, _>(0..SMALL_SIZE as u16);
3562    }
3563
3564    #[test]
3565    fn u32_single_column() {
3566        required_and_optional::<UInt32Array, _>(0..SMALL_SIZE as u32);
3567    }
3568
3569    #[test]
3570    fn u64_single_column() {
3571        required_and_optional::<UInt64Array, _>(0..SMALL_SIZE as u64);
3572    }
3573
3574    #[test]
3575    fn f32_single_column() {
3576        required_and_optional::<Float32Array, _>((0..SMALL_SIZE).map(|i| i as f32));
3577    }
3578
3579    #[test]
3580    fn f64_single_column() {
3581        required_and_optional::<Float64Array, _>((0..SMALL_SIZE).map(|i| i as f64));
3582    }
3583
3584    // The timestamp array types don't implement From<Vec<T>> because they need the timezone
3585    // argument, and they also doesn't support building from a Vec<Option<T>>, so call
3586    // RoundTripTest manually instead of calling required_and_optional for these tests.
3587
3588    #[test]
3589    fn timestamp_second_single_column() {
3590        let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
3591        let values = Arc::new(TimestampSecondArray::from(raw_values));
3592
3593        RoundTripTest::new(values).with_nullable(false).run();
3594    }
3595
3596    #[test]
3597    fn timestamp_millisecond_single_column() {
3598        let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
3599        let values = Arc::new(TimestampMillisecondArray::from(raw_values));
3600
3601        RoundTripTest::new(values).with_nullable(false).run();
3602    }
3603
3604    #[test]
3605    fn timestamp_microsecond_single_column() {
3606        let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
3607        let values = Arc::new(TimestampMicrosecondArray::from(raw_values));
3608
3609        RoundTripTest::new(values).with_nullable(false).run();
3610    }
3611
3612    #[test]
3613    fn timestamp_nanosecond_single_column() {
3614        let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
3615        let values = Arc::new(TimestampNanosecondArray::from(raw_values));
3616
3617        RoundTripTest::new(values).with_nullable(false).run();
3618    }
3619
3620    #[test]
3621    fn date32_single_column() {
3622        required_and_optional::<Date32Array, _>(0..SMALL_SIZE as i32);
3623    }
3624
3625    #[test]
3626    fn date64_single_column() {
3627        // Date64 must be a multiple of 86400000, see ARROW-10925
3628        required_and_optional::<Date64Array, _>(
3629            (0..(SMALL_SIZE as i64 * 86400000)).step_by(86400000),
3630        );
3631    }
3632
3633    #[test]
3634    fn time32_second_single_column() {
3635        required_and_optional::<Time32SecondArray, _>(0..SMALL_SIZE as i32);
3636    }
3637
3638    #[test]
3639    fn time32_millisecond_single_column() {
3640        required_and_optional::<Time32MillisecondArray, _>(0..SMALL_SIZE as i32);
3641    }
3642
3643    #[test]
3644    fn time64_microsecond_single_column() {
3645        required_and_optional::<Time64MicrosecondArray, _>(0..SMALL_SIZE as i64);
3646    }
3647
3648    #[test]
3649    fn time64_nanosecond_single_column() {
3650        required_and_optional::<Time64NanosecondArray, _>(0..SMALL_SIZE as i64);
3651    }
3652
3653    #[test]
3654    fn duration_second_single_column() {
3655        required_and_optional::<DurationSecondArray, _>(0..SMALL_SIZE as i64);
3656    }
3657
3658    #[test]
3659    fn duration_millisecond_single_column() {
3660        required_and_optional::<DurationMillisecondArray, _>(0..SMALL_SIZE as i64);
3661    }
3662
3663    #[test]
3664    fn duration_microsecond_single_column() {
3665        required_and_optional::<DurationMicrosecondArray, _>(0..SMALL_SIZE as i64);
3666    }
3667
3668    #[test]
3669    fn duration_nanosecond_single_column() {
3670        required_and_optional::<DurationNanosecondArray, _>(0..SMALL_SIZE as i64);
3671    }
3672
3673    #[test]
3674    fn interval_year_month_single_column() {
3675        required_and_optional::<IntervalYearMonthArray, _>(0..SMALL_SIZE as i32);
3676    }
3677
3678    #[test]
3679    fn interval_day_time_single_column() {
3680        required_and_optional::<IntervalDayTimeArray, _>(vec![
3681            IntervalDayTime::new(0, 1),
3682            IntervalDayTime::new(0, 3),
3683            IntervalDayTime::new(3, -2),
3684            IntervalDayTime::new(-200, 4),
3685        ]);
3686    }
3687
3688    #[test]
3689    #[should_panic(
3690        expected = "Attempting to write an Arrow interval type MonthDayNano to parquet that is not yet implemented"
3691    )]
3692    fn interval_month_day_nano_single_column() {
3693        required_and_optional::<IntervalMonthDayNanoArray, _>(vec![
3694            IntervalMonthDayNano::new(0, 1, 5),
3695            IntervalMonthDayNano::new(0, 3, 2),
3696            IntervalMonthDayNano::new(3, -2, -5),
3697            IntervalMonthDayNano::new(-200, 4, -1),
3698        ]);
3699    }
3700
3701    #[test]
3702    fn binary_single_column() {
3703        let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
3704        let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
3705        let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
3706
3707        // BinaryArrays can't be built from Vec<Option<&str>>, so only call `values_required`
3708        values_required::<BinaryArray, _>(many_vecs_iter);
3709    }
3710
3711    #[test]
3712    fn binary_view_single_column() {
3713        let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
3714        let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
3715        let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
3716
3717        // BinaryArrays can't be built from Vec<Option<&str>>, so only call `values_required`
3718        values_required::<BinaryViewArray, _>(many_vecs_iter);
3719    }
3720
3721    #[test]
3722    fn i32_column_bloom_filter_at_end() {
3723        let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
3724        let files = RoundTripTest::new(array)
3725            .with_nullable(false)
3726            .with_bloom_filter(true)
3727            .with_bloom_filter_position(BloomFilterPosition::End)
3728            .run();
3729
3730        check_bloom_filter(
3731            files,
3732            "col".to_string(),
3733            (0..SMALL_SIZE as i32).collect(),
3734            (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
3735        );
3736    }
3737
3738    #[test]
3739    fn i32_column_bloom_filter() {
3740        let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
3741        let files = RoundTripTest::new(array)
3742            .with_nullable(false)
3743            .with_bloom_filter(true)
3744            .run();
3745
3746        check_bloom_filter(
3747            files,
3748            "col".to_string(),
3749            (0..SMALL_SIZE as i32).collect(),
3750            (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
3751        );
3752    }
3753
3754    /// Test that bloom filter folding produces correct results even when
3755    /// the configured NDV differs significantly from actual NDV.
3756    /// A large NDV means a larger initial filter that gets folded down;
3757    /// a small NDV means a smaller initial filter.
3758    #[test]
3759    fn i32_column_bloom_filter_fixed_ndv() {
3760        let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
3761
3762        // NDV much larger than actual distinct values — tests folding a large filter down
3763        let files = RoundTripTest::new(array.clone())
3764            .with_nullable(false)
3765            .with_bloom_filter(true)
3766            .with_bloom_filter_ndv(1_000_000)
3767            .run();
3768
3769        check_bloom_filter(
3770            files,
3771            "col".to_string(),
3772            (0..SMALL_SIZE as i32).collect(),
3773            (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
3774        );
3775
3776        // NDV smaller than actual distinct values — tests the underestimate path
3777        let files = RoundTripTest::new(array)
3778            .with_nullable(false)
3779            .with_bloom_filter(true)
3780            .with_bloom_filter_ndv(3)
3781            .run();
3782
3783        check_bloom_filter(
3784            files,
3785            "col".to_string(),
3786            (0..SMALL_SIZE as i32).collect(),
3787            (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
3788        );
3789    }
3790
3791    #[test]
3792    fn binary_column_bloom_filter() {
3793        let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
3794        let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
3795        let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
3796
3797        let array = Arc::new(BinaryArray::from_iter_values(many_vecs_iter));
3798        let files = RoundTripTest::new(array)
3799            .with_nullable(false)
3800            .with_bloom_filter(true)
3801            .run();
3802
3803        check_bloom_filter(
3804            files,
3805            "col".to_string(),
3806            many_vecs,
3807            vec![vec![(SMALL_SIZE + 1) as u8]],
3808        );
3809    }
3810
3811    #[test]
3812    fn empty_string_null_column_bloom_filter() {
3813        let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
3814        let raw_strs = raw_values.iter().map(|s| s.as_str());
3815
3816        let array = Arc::new(StringArray::from_iter_values(raw_strs));
3817        let files = RoundTripTest::new(array)
3818            .with_nullable(false)
3819            .with_bloom_filter(true)
3820            .run();
3821
3822        let optional_raw_values: Vec<_> = raw_values
3823            .iter()
3824            .enumerate()
3825            .filter_map(|(i, v)| if i % 2 == 0 { None } else { Some(v.as_str()) })
3826            .collect();
3827        // For null slots, empty string should not be in bloom filter.
3828        check_bloom_filter(files, "col".to_string(), optional_raw_values, vec![""]);
3829    }
3830
3831    #[test]
3832    fn large_binary_single_column() {
3833        let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
3834        let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
3835        let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
3836
3837        // LargeBinaryArrays can't be built from Vec<Option<&str>>, so only call `values_required`
3838        values_required::<LargeBinaryArray, _>(many_vecs_iter);
3839    }
3840
3841    #[test]
3842    fn fixed_size_binary_single_column() {
3843        let mut builder = FixedSizeBinaryBuilder::new(4);
3844        builder.append_value(b"0123").unwrap();
3845        builder.append_null();
3846        builder.append_value(b"8910").unwrap();
3847        builder.append_value(b"1112").unwrap();
3848        let array = Arc::new(builder.finish());
3849
3850        RoundTripTest::new(array).run();
3851    }
3852
3853    #[test]
3854    fn string_single_column() {
3855        let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
3856        let raw_strs = raw_values.iter().map(|s| s.as_str());
3857
3858        required_and_optional::<StringArray, _>(raw_strs);
3859    }
3860
3861    #[test]
3862    fn large_string_single_column() {
3863        let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
3864        let raw_strs = raw_values.iter().map(|s| s.as_str());
3865
3866        required_and_optional::<LargeStringArray, _>(raw_strs);
3867    }
3868
3869    #[test]
3870    fn string_view_single_column() {
3871        let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
3872        let raw_strs = raw_values.iter().map(|s| s.as_str());
3873
3874        required_and_optional::<StringViewArray, _>(raw_strs);
3875    }
3876
3877    #[test]
3878    fn null_list_single_column() {
3879        let null_field = Field::new_list_field(DataType::Null, true);
3880        let list_field = Field::new("emptylist", DataType::List(Arc::new(null_field)), true);
3881
3882        let schema = Schema::new(vec![list_field]);
3883
3884        // Build [[], null, [null, null]]
3885        let a_values = NullArray::new(2);
3886        let a_value_offsets = arrow::buffer::Buffer::from([0, 0, 0, 2].to_byte_slice());
3887        let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
3888            DataType::Null,
3889            true,
3890        ))))
3891        .len(3)
3892        .add_buffer(a_value_offsets)
3893        .null_bit_buffer(Some(Buffer::from([0b00000101])))
3894        .add_child_data(a_values.into_data())
3895        .build()
3896        .unwrap();
3897
3898        let a = ListArray::from(a_list_data);
3899
3900        assert!(a.is_valid(0));
3901        assert!(!a.is_valid(1));
3902        assert!(a.is_valid(2));
3903
3904        assert_eq!(a.value(0).len(), 0);
3905        assert_eq!(a.value(2).len(), 2);
3906        assert_eq!(a.value(2).logical_nulls().unwrap().null_count(), 2);
3907
3908        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
3909        roundtrip(batch, None);
3910    }
3911
3912    #[test]
3913    fn list_single_column() {
3914        let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
3915        let a_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice());
3916        let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
3917            DataType::Int32,
3918            false,
3919        ))))
3920        .len(5)
3921        .add_buffer(a_value_offsets)
3922        .null_bit_buffer(Some(Buffer::from([0b00011011])))
3923        .add_child_data(a_values.into_data())
3924        .build()
3925        .unwrap();
3926
3927        assert_eq!(a_list_data.null_count(), 1);
3928
3929        let a = ListArray::from(a_list_data);
3930        let values = Arc::new(a);
3931
3932        RoundTripTest::new(values).run();
3933    }
3934
3935    #[test]
3936    fn large_list_single_column() {
3937        let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
3938        let a_value_offsets = arrow::buffer::Buffer::from([0i64, 1, 3, 3, 6, 10].to_byte_slice());
3939        let a_list_data = ArrayData::builder(DataType::LargeList(Arc::new(Field::new(
3940            "large_item",
3941            DataType::Int32,
3942            true,
3943        ))))
3944        .len(5)
3945        .add_buffer(a_value_offsets)
3946        .add_child_data(a_values.into_data())
3947        .null_bit_buffer(Some(Buffer::from([0b00011011])))
3948        .build()
3949        .unwrap();
3950
3951        // I think this setup is incorrect because this should pass
3952        assert_eq!(a_list_data.null_count(), 1);
3953
3954        let a = LargeListArray::from(a_list_data);
3955        let values = Arc::new(a);
3956
3957        RoundTripTest::new(values).run();
3958    }
3959
3960    #[test]
3961    fn list_nested_nulls() {
3962        use arrow::datatypes::Int32Type;
3963        let data = vec![
3964            Some(vec![Some(1)]),
3965            Some(vec![Some(2), Some(3)]),
3966            None,
3967            Some(vec![Some(4), Some(5), None]),
3968            Some(vec![None]),
3969            Some(vec![Some(6), Some(7)]),
3970        ];
3971
3972        let list = ListArray::from_iter_primitive::<Int32Type, _, _>(data.clone());
3973        RoundTripTest::new(Arc::new(list)).run();
3974
3975        let list = LargeListArray::from_iter_primitive::<Int32Type, _, _>(data);
3976        RoundTripTest::new(Arc::new(list)).run();
3977    }
3978
3979    #[test]
3980    fn list_utf8_view_selective_padding_roundtrip() {
3981        let item = Arc::new(Field::new_list_field(DataType::Utf8View, true));
3982        let mut builder = ListBuilder::new(StringViewBuilder::new()).with_field(item);
3983        builder.values().append_value("a");
3984        builder.values().append_null();
3985        builder.append(true);
3986        // The null parent list covers selective padding dropping values below
3987        // the list definition level while preserving the preceding item null.
3988        builder.append(false);
3989        // The long string covers the non-inlined Utf8View buffer path.
3990        builder.values().append_value("large payload over 12 bytes");
3991        builder.append(true);
3992
3993        RoundTripTest::new(Arc::new(builder.finish())).run();
3994    }
3995
3996    #[test]
3997    fn struct_single_column() {
3998        let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
3999        let struct_field_a = Arc::new(Field::new("f", DataType::Int32, false));
4000        let s = StructArray::from(vec![(struct_field_a, Arc::new(a_values) as ArrayRef)]);
4001
4002        let values = Arc::new(s);
4003        RoundTripTest::new(values).with_nullable(false).run();
4004    }
4005
4006    #[test]
4007    fn list_and_map_coerced_names() {
4008        // Create map and list with non-Parquet naming
4009        let list_field =
4010            Field::new_list("my_list", Field::new("item", DataType::Int32, false), false);
4011        let map_field = Field::new_map(
4012            "my_map",
4013            "my_entries",
4014            Field::new("my_keys", DataType::Int32, false),
4015            Field::new("my_values", DataType::Int32, true),
4016            false,
4017            true,
4018        );
4019
4020        let list_array = create_random_array(&list_field, 100, 0.0, 0.0).unwrap();
4021        let map_array = create_random_array(&map_field, 100, 0.0, 0.0).unwrap();
4022
4023        let arrow_schema = Arc::new(Schema::new(vec![list_field, map_field]));
4024
4025        // Write data to Parquet but coerce names to match spec
4026        let props = Some(WriterProperties::builder().set_coerce_types(true).build());
4027        let file = tempfile::tempfile().unwrap();
4028        let mut writer =
4029            ArrowWriter::try_new(file.try_clone().unwrap(), arrow_schema.clone(), props).unwrap();
4030
4031        let batch = RecordBatch::try_new(arrow_schema, vec![list_array, map_array]).unwrap();
4032        writer.write(&batch).unwrap();
4033        let file_metadata = writer.close().unwrap();
4034
4035        let schema = file_metadata.file_metadata().schema();
4036        // Coerced name of "item" should be "element"
4037        let list_field = &schema.get_fields()[0].get_fields()[0];
4038        assert_eq!(list_field.get_fields()[0].name(), "element");
4039
4040        let map_field = &schema.get_fields()[1].get_fields()[0];
4041        // Coerced name of "entries" should be "key_value"
4042        assert_eq!(map_field.name(), "key_value");
4043        // Coerced name of "my_keys" should be "key"
4044        assert_eq!(map_field.get_fields()[0].name(), "key");
4045        // Coerced name of "my_values" should be "value"
4046        assert_eq!(map_field.get_fields()[1].name(), "value");
4047
4048        // Double check schema after reading from the file
4049        let reader = SerializedFileReader::new(file).unwrap();
4050        let file_schema = reader.metadata().file_metadata().schema();
4051        let fields = file_schema.get_fields();
4052        let list_field = &fields[0].get_fields()[0];
4053        assert_eq!(list_field.get_fields()[0].name(), "element");
4054        let map_field = &fields[1].get_fields()[0];
4055        assert_eq!(map_field.name(), "key_value");
4056        assert_eq!(map_field.get_fields()[0].name(), "key");
4057        assert_eq!(map_field.get_fields()[1].name(), "value");
4058    }
4059
4060    #[test]
4061    fn fallback_flush_data_page() {
4062        //tests if the Fallback::flush_data_page clears all buffers correctly
4063        let raw_values: Vec<_> = (0..MEDIUM_SIZE).map(|i| i.to_string()).collect();
4064        let values = Arc::new(StringArray::from(raw_values));
4065        let encodings = vec![
4066            Encoding::DELTA_BYTE_ARRAY,
4067            Encoding::DELTA_LENGTH_BYTE_ARRAY,
4068        ];
4069        let data_type = values.data_type().clone();
4070        let schema = Arc::new(Schema::new(vec![Field::new("col", data_type, false)]));
4071        let expected_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4072
4073        let row_group_sizes = [1024, SMALL_SIZE, SMALL_SIZE / 2, SMALL_SIZE / 2 + 1, 10];
4074        let data_page_size_limit: usize = 32;
4075        let write_batch_size: usize = 16;
4076
4077        for encoding in &encodings {
4078            for row_group_size in row_group_sizes {
4079                let props = WriterProperties::builder()
4080                    .set_writer_version(WriterVersion::PARQUET_2_0)
4081                    .set_max_row_group_row_count(Some(row_group_size))
4082                    .set_dictionary_enabled(false)
4083                    .set_encoding(*encoding)
4084                    .set_data_page_size_limit(data_page_size_limit)
4085                    .set_write_batch_size(write_batch_size)
4086                    .build();
4087
4088                roundtrip_opts_with_array_validation(&expected_batch, props, |a, b| {
4089                    let string_array_a = StringArray::from(a.clone());
4090                    let string_array_b = StringArray::from(b.clone());
4091                    let vec_a: Vec<&str> = string_array_a.iter().map(|v| v.unwrap()).collect();
4092                    let vec_b: Vec<&str> = string_array_b.iter().map(|v| v.unwrap()).collect();
4093                    assert_eq!(
4094                        vec_a, vec_b,
4095                        "failed for encoder: {encoding:?} and row_group_size: {row_group_size:?}"
4096                    );
4097                });
4098            }
4099        }
4100    }
4101
4102    #[test]
4103    fn arrow_writer_string_dictionary() {
4104        // define schema
4105        #[allow(deprecated)]
4106        let schema = Arc::new(Schema::new(vec![Field::new_dict(
4107            "dictionary",
4108            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4109            true,
4110            42,
4111            true,
4112        )]));
4113
4114        // create some data
4115        let d: Int32DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")]
4116            .iter()
4117            .copied()
4118            .collect();
4119
4120        // build a record batch
4121        RoundTripTest::new(Arc::new(d)).with_schema(schema).run();
4122    }
4123
4124    #[test]
4125    fn arrow_writer_test_type_compatibility() {
4126        fn ensure_compatible_write<T1, T2>(array1: T1, array2: T2, expected_result: T1)
4127        where
4128            T1: Array + 'static,
4129            T2: Array + 'static,
4130        {
4131            let schema1 = Arc::new(Schema::new(vec![Field::new(
4132                "a",
4133                array1.data_type().clone(),
4134                false,
4135            )]));
4136
4137            let file = tempfile().unwrap();
4138            let mut writer =
4139                ArrowWriter::try_new(file.try_clone().unwrap(), schema1.clone(), None).unwrap();
4140
4141            let rb1 = RecordBatch::try_new(schema1.clone(), vec![Arc::new(array1)]).unwrap();
4142            writer.write(&rb1).unwrap();
4143
4144            let schema2 = Arc::new(Schema::new(vec![Field::new(
4145                "a",
4146                array2.data_type().clone(),
4147                false,
4148            )]));
4149            let rb2 = RecordBatch::try_new(schema2, vec![Arc::new(array2)]).unwrap();
4150            writer.write(&rb2).unwrap();
4151
4152            writer.close().unwrap();
4153
4154            let mut record_batch_reader =
4155                ParquetRecordBatchReader::try_new(file.try_clone().unwrap(), 1024).unwrap();
4156            let actual_batch = record_batch_reader.next().unwrap().unwrap();
4157
4158            let expected_batch =
4159                RecordBatch::try_new(schema1, vec![Arc::new(expected_result)]).unwrap();
4160            assert_eq!(actual_batch, expected_batch);
4161        }
4162
4163        // check compatibility between native and dictionaries
4164
4165        ensure_compatible_write(
4166            DictionaryArray::new(
4167                UInt8Array::from_iter_values(vec![0]),
4168                Arc::new(StringArray::from_iter_values(vec!["parquet"])),
4169            ),
4170            StringArray::from_iter_values(vec!["barquet"]),
4171            DictionaryArray::new(
4172                UInt8Array::from_iter_values(vec![0, 1]),
4173                Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])),
4174            ),
4175        );
4176
4177        ensure_compatible_write(
4178            StringArray::from_iter_values(vec!["parquet"]),
4179            DictionaryArray::new(
4180                UInt8Array::from_iter_values(vec![0]),
4181                Arc::new(StringArray::from_iter_values(vec!["barquet"])),
4182            ),
4183            StringArray::from_iter_values(vec!["parquet", "barquet"]),
4184        );
4185
4186        // check compatibility between dictionaries with different key types
4187
4188        ensure_compatible_write(
4189            DictionaryArray::new(
4190                UInt8Array::from_iter_values(vec![0]),
4191                Arc::new(StringArray::from_iter_values(vec!["parquet"])),
4192            ),
4193            DictionaryArray::new(
4194                UInt16Array::from_iter_values(vec![0]),
4195                Arc::new(StringArray::from_iter_values(vec!["barquet"])),
4196            ),
4197            DictionaryArray::new(
4198                UInt8Array::from_iter_values(vec![0, 1]),
4199                Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])),
4200            ),
4201        );
4202
4203        // check compatibility between dictionaries with different value types
4204        ensure_compatible_write(
4205            DictionaryArray::new(
4206                UInt8Array::from_iter_values(vec![0]),
4207                Arc::new(StringArray::from_iter_values(vec!["parquet"])),
4208            ),
4209            DictionaryArray::new(
4210                UInt8Array::from_iter_values(vec![0]),
4211                Arc::new(LargeStringArray::from_iter_values(vec!["barquet"])),
4212            ),
4213            DictionaryArray::new(
4214                UInt8Array::from_iter_values(vec![0, 1]),
4215                Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])),
4216            ),
4217        );
4218
4219        // check compatibility between a dictionary and a native array with a different type
4220        ensure_compatible_write(
4221            DictionaryArray::new(
4222                UInt8Array::from_iter_values(vec![0]),
4223                Arc::new(StringArray::from_iter_values(vec!["parquet"])),
4224            ),
4225            LargeStringArray::from_iter_values(vec!["barquet"]),
4226            DictionaryArray::new(
4227                UInt8Array::from_iter_values(vec![0, 1]),
4228                Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])),
4229            ),
4230        );
4231
4232        // check compatibility for string types
4233
4234        ensure_compatible_write(
4235            StringArray::from_iter_values(vec!["parquet"]),
4236            LargeStringArray::from_iter_values(vec!["barquet"]),
4237            StringArray::from_iter_values(vec!["parquet", "barquet"]),
4238        );
4239
4240        ensure_compatible_write(
4241            LargeStringArray::from_iter_values(vec!["parquet"]),
4242            StringArray::from_iter_values(vec!["barquet"]),
4243            LargeStringArray::from_iter_values(vec!["parquet", "barquet"]),
4244        );
4245
4246        ensure_compatible_write(
4247            StringArray::from_iter_values(vec!["parquet"]),
4248            StringViewArray::from_iter_values(vec!["barquet"]),
4249            StringArray::from_iter_values(vec!["parquet", "barquet"]),
4250        );
4251
4252        ensure_compatible_write(
4253            StringViewArray::from_iter_values(vec!["parquet"]),
4254            StringArray::from_iter_values(vec!["barquet"]),
4255            StringViewArray::from_iter_values(vec!["parquet", "barquet"]),
4256        );
4257
4258        ensure_compatible_write(
4259            LargeStringArray::from_iter_values(vec!["parquet"]),
4260            StringViewArray::from_iter_values(vec!["barquet"]),
4261            LargeStringArray::from_iter_values(vec!["parquet", "barquet"]),
4262        );
4263
4264        ensure_compatible_write(
4265            StringViewArray::from_iter_values(vec!["parquet"]),
4266            LargeStringArray::from_iter_values(vec!["barquet"]),
4267            StringViewArray::from_iter_values(vec!["parquet", "barquet"]),
4268        );
4269
4270        // check compatibility for binary types
4271
4272        ensure_compatible_write(
4273            BinaryArray::from_iter_values(vec![b"parquet"]),
4274            LargeBinaryArray::from_iter_values(vec![b"barquet"]),
4275            BinaryArray::from_iter_values(vec![b"parquet", b"barquet"]),
4276        );
4277
4278        ensure_compatible_write(
4279            LargeBinaryArray::from_iter_values(vec![b"parquet"]),
4280            BinaryArray::from_iter_values(vec![b"barquet"]),
4281            LargeBinaryArray::from_iter_values(vec![b"parquet", b"barquet"]),
4282        );
4283
4284        ensure_compatible_write(
4285            BinaryArray::from_iter_values(vec![b"parquet"]),
4286            BinaryViewArray::from_iter_values(vec![b"barquet"]),
4287            BinaryArray::from_iter_values(vec![b"parquet", b"barquet"]),
4288        );
4289
4290        ensure_compatible_write(
4291            BinaryViewArray::from_iter_values(vec![b"parquet"]),
4292            BinaryArray::from_iter_values(vec![b"barquet"]),
4293            BinaryViewArray::from_iter_values(vec![b"parquet", b"barquet"]),
4294        );
4295
4296        ensure_compatible_write(
4297            BinaryViewArray::from_iter_values(vec![b"parquet"]),
4298            LargeBinaryArray::from_iter_values(vec![b"barquet"]),
4299            BinaryViewArray::from_iter_values(vec![b"parquet", b"barquet"]),
4300        );
4301
4302        ensure_compatible_write(
4303            LargeBinaryArray::from_iter_values(vec![b"parquet"]),
4304            BinaryViewArray::from_iter_values(vec![b"barquet"]),
4305            LargeBinaryArray::from_iter_values(vec![b"parquet", b"barquet"]),
4306        );
4307
4308        // check compatibility for list types
4309
4310        let list_field_metadata = HashMap::from_iter(vec![(
4311            PARQUET_FIELD_ID_META_KEY.to_string(),
4312            "1".to_string(),
4313        )]);
4314        let list_field = Field::new_list_field(DataType::Int32, false);
4315
4316        let values1 = Arc::new(Int32Array::from(vec![0, 1, 2, 3, 4]));
4317        let offsets1 = OffsetBuffer::new(vec![0, 2, 5].into());
4318
4319        let values2 = Arc::new(Int32Array::from(vec![5, 6, 7, 8, 9]));
4320        let offsets2 = OffsetBuffer::new(vec![0, 3, 5].into());
4321
4322        let values_expected = Arc::new(Int32Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]));
4323        let offsets_expected = OffsetBuffer::new(vec![0, 2, 5, 8, 10].into());
4324
4325        ensure_compatible_write(
4326            // when the initial schema has the metadata ...
4327            ListArray::try_new(
4328                Arc::new(
4329                    list_field
4330                        .clone()
4331                        .with_metadata(list_field_metadata.clone()),
4332                ),
4333                offsets1,
4334                values1,
4335                None,
4336            )
4337            .unwrap(),
4338            // ... and some intermediate schema doesn't have the metadata
4339            ListArray::try_new(Arc::new(list_field.clone()), offsets2, values2, None).unwrap(),
4340            // ... the write will still go through, and the resulting schema will inherit the initial metadata
4341            ListArray::try_new(
4342                Arc::new(
4343                    list_field
4344                        .clone()
4345                        .with_metadata(list_field_metadata.clone()),
4346                ),
4347                offsets_expected,
4348                values_expected,
4349                None,
4350            )
4351            .unwrap(),
4352        );
4353    }
4354
4355    #[test]
4356    fn arrow_writer_primitive_dictionary() {
4357        // define schema
4358        #[allow(deprecated)]
4359        let schema = Arc::new(Schema::new(vec![Field::new_dict(
4360            "dictionary",
4361            DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::UInt32)),
4362            true,
4363            42,
4364            true,
4365        )]));
4366
4367        // create some data
4368        let mut builder = PrimitiveDictionaryBuilder::<UInt8Type, UInt32Type>::new();
4369        builder.append(12345678).unwrap();
4370        builder.append_null();
4371        builder.append(22345678).unwrap();
4372        builder.append(12345678).unwrap();
4373        let d = builder.finish();
4374
4375        RoundTripTest::new(Arc::new(d)).with_schema(schema).run();
4376    }
4377
4378    #[test]
4379    fn arrow_writer_decimal32_dictionary() {
4380        let integers = vec![12345, 56789, 34567];
4381
4382        let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]);
4383
4384        let values = Decimal32Array::from(integers.clone())
4385            .with_precision_and_scale(5, 2)
4386            .unwrap();
4387
4388        let array = DictionaryArray::new(keys, Arc::new(values));
4389        RoundTripTest::new(Arc::new(array.clone())).run();
4390
4391        let values = Decimal32Array::from(integers)
4392            .with_precision_and_scale(9, 2)
4393            .unwrap();
4394
4395        let array = array.with_values(Arc::new(values));
4396        RoundTripTest::new(Arc::new(array)).run();
4397    }
4398
4399    #[test]
4400    fn arrow_writer_decimal64_dictionary() {
4401        let integers = vec![12345, 56789, 34567];
4402
4403        let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]);
4404
4405        let values = Decimal64Array::from(integers.clone())
4406            .with_precision_and_scale(5, 2)
4407            .unwrap();
4408
4409        let array = DictionaryArray::new(keys, Arc::new(values));
4410        RoundTripTest::new(Arc::new(array.clone())).run();
4411
4412        let values = Decimal64Array::from(integers)
4413            .with_precision_and_scale(12, 2)
4414            .unwrap();
4415
4416        let array = array.with_values(Arc::new(values));
4417        RoundTripTest::new(Arc::new(array)).run();
4418    }
4419
4420    #[test]
4421    fn arrow_writer_decimal128_dictionary() {
4422        let integers = vec![12345, 56789, 34567];
4423
4424        let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]);
4425
4426        let values = Decimal128Array::from(integers.clone())
4427            .with_precision_and_scale(5, 2)
4428            .unwrap();
4429
4430        let array = DictionaryArray::new(keys, Arc::new(values));
4431        RoundTripTest::new(Arc::new(array.clone())).run();
4432
4433        let values = Decimal128Array::from(integers)
4434            .with_precision_and_scale(12, 2)
4435            .unwrap();
4436
4437        let array = array.with_values(Arc::new(values));
4438        RoundTripTest::new(Arc::new(array)).run();
4439    }
4440
4441    #[test]
4442    fn arrow_writer_decimal256_dictionary() {
4443        let integers = vec![
4444            i256::from_i128(12345),
4445            i256::from_i128(56789),
4446            i256::from_i128(34567),
4447        ];
4448
4449        let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]);
4450
4451        let values = Decimal256Array::from(integers.clone())
4452            .with_precision_and_scale(5, 2)
4453            .unwrap();
4454
4455        let array = DictionaryArray::new(keys, Arc::new(values));
4456        RoundTripTest::new(Arc::new(array.clone())).run();
4457
4458        let values = Decimal256Array::from(integers)
4459            .with_precision_and_scale(12, 2)
4460            .unwrap();
4461
4462        let array = array.with_values(Arc::new(values));
4463        RoundTripTest::new(Arc::new(array)).run();
4464    }
4465
4466    #[test]
4467    fn arrow_writer_string_dictionary_unsigned_index() {
4468        // define schema
4469        #[allow(deprecated)]
4470        let schema = Arc::new(Schema::new(vec![Field::new_dict(
4471            "dictionary",
4472            DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
4473            true,
4474            42,
4475            true,
4476        )]));
4477
4478        // create some data
4479        let d: UInt8DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")]
4480            .iter()
4481            .copied()
4482            .collect();
4483
4484        RoundTripTest::new(Arc::new(d)).with_schema(schema).run();
4485    }
4486
4487    #[test]
4488    fn u32_min_max() {
4489        // check values roundtrip through parquet
4490        let src = [
4491            u32::MIN,
4492            1,
4493            (i32::MAX as u32) - 1,
4494            i32::MAX as u32,
4495            (i32::MAX as u32) + 1,
4496            u32::MAX - 1,
4497            u32::MAX,
4498        ];
4499        let values = Arc::new(UInt32Array::from_iter_values(src.iter().cloned()));
4500        let files = RoundTripTest::new(values).with_nullable(false).run();
4501
4502        for file in files {
4503            // check statistics are valid
4504            let reader = SerializedFileReader::new(file).unwrap();
4505            let metadata = reader.metadata();
4506
4507            let mut row_offset = 0;
4508            for row_group in metadata.row_groups() {
4509                assert_eq!(row_group.num_columns(), 1);
4510                let column = row_group.column(0);
4511
4512                let num_values = column.num_values() as usize;
4513                let src_slice = &src[row_offset..row_offset + num_values];
4514                row_offset += column.num_values() as usize;
4515
4516                let stats = column.statistics().unwrap();
4517                if let Statistics::Int32(stats) = stats {
4518                    assert_eq!(
4519                        *stats.min_opt().unwrap() as u32,
4520                        *src_slice.iter().min().unwrap()
4521                    );
4522                    assert_eq!(
4523                        *stats.max_opt().unwrap() as u32,
4524                        *src_slice.iter().max().unwrap()
4525                    );
4526                } else {
4527                    panic!("Statistics::Int32 missing")
4528                }
4529            }
4530        }
4531    }
4532
4533    #[test]
4534    fn u64_min_max() {
4535        // check values roundtrip through parquet
4536        let src = [
4537            u64::MIN,
4538            1,
4539            (i64::MAX as u64) - 1,
4540            i64::MAX as u64,
4541            (i64::MAX as u64) + 1,
4542            u64::MAX - 1,
4543            u64::MAX,
4544        ];
4545        let values = Arc::new(UInt64Array::from_iter_values(src.iter().cloned()));
4546        let files = RoundTripTest::new(values).with_nullable(false).run();
4547
4548        for file in files {
4549            // check statistics are valid
4550            let reader = SerializedFileReader::new(file).unwrap();
4551            let metadata = reader.metadata();
4552
4553            let mut row_offset = 0;
4554            for row_group in metadata.row_groups() {
4555                assert_eq!(row_group.num_columns(), 1);
4556                let column = row_group.column(0);
4557
4558                let num_values = column.num_values() as usize;
4559                let src_slice = &src[row_offset..row_offset + num_values];
4560                row_offset += column.num_values() as usize;
4561
4562                let stats = column.statistics().unwrap();
4563                if let Statistics::Int64(stats) = stats {
4564                    assert_eq!(
4565                        *stats.min_opt().unwrap() as u64,
4566                        *src_slice.iter().min().unwrap()
4567                    );
4568                    assert_eq!(
4569                        *stats.max_opt().unwrap() as u64,
4570                        *src_slice.iter().max().unwrap()
4571                    );
4572                } else {
4573                    panic!("Statistics::Int64 missing")
4574                }
4575            }
4576        }
4577    }
4578
4579    #[test]
4580    fn statistics_null_counts_only_nulls() {
4581        // check that null-count statistics for "only NULL"-columns are correct
4582        let values = Arc::new(UInt64Array::from(vec![None, None]));
4583        let files = RoundTripTest::new(values).run();
4584
4585        for file in files {
4586            // check statistics are valid
4587            let reader = SerializedFileReader::new(file).unwrap();
4588            let metadata = reader.metadata();
4589            assert_eq!(metadata.num_row_groups(), 1);
4590            let row_group = metadata.row_group(0);
4591            assert_eq!(row_group.num_columns(), 1);
4592            let column = row_group.column(0);
4593            let stats = column.statistics().unwrap();
4594            assert_eq!(stats.null_count_opt(), Some(2));
4595        }
4596    }
4597
4598    #[test]
4599    fn test_list_of_struct_roundtrip() {
4600        // define schema
4601        let int_field = Field::new("a", DataType::Int32, true);
4602        let int_field2 = Field::new("b", DataType::Int32, true);
4603
4604        let int_builder = Int32Builder::with_capacity(10);
4605        let int_builder2 = Int32Builder::with_capacity(10);
4606
4607        let struct_builder = StructBuilder::new(
4608            vec![int_field, int_field2],
4609            vec![Box::new(int_builder), Box::new(int_builder2)],
4610        );
4611        let mut list_builder = ListBuilder::new(struct_builder);
4612
4613        // Construct the following array
4614        // [{a: 1, b: 2}], [], null, [null, null], [{a: null, b: 3}], [{a: 2, b: null}]
4615
4616        // [{a: 1, b: 2}]
4617        let values = list_builder.values();
4618        values
4619            .field_builder::<Int32Builder>(0)
4620            .unwrap()
4621            .append_value(1);
4622        values
4623            .field_builder::<Int32Builder>(1)
4624            .unwrap()
4625            .append_value(2);
4626        values.append(true);
4627        list_builder.append(true);
4628
4629        // []
4630        list_builder.append(true);
4631
4632        // null
4633        list_builder.append(false);
4634
4635        // [null, null]
4636        let values = list_builder.values();
4637        values
4638            .field_builder::<Int32Builder>(0)
4639            .unwrap()
4640            .append_null();
4641        values
4642            .field_builder::<Int32Builder>(1)
4643            .unwrap()
4644            .append_null();
4645        values.append(false);
4646        values
4647            .field_builder::<Int32Builder>(0)
4648            .unwrap()
4649            .append_null();
4650        values
4651            .field_builder::<Int32Builder>(1)
4652            .unwrap()
4653            .append_null();
4654        values.append(false);
4655        list_builder.append(true);
4656
4657        // [{a: null, b: 3}]
4658        let values = list_builder.values();
4659        values
4660            .field_builder::<Int32Builder>(0)
4661            .unwrap()
4662            .append_null();
4663        values
4664            .field_builder::<Int32Builder>(1)
4665            .unwrap()
4666            .append_value(3);
4667        values.append(true);
4668        list_builder.append(true);
4669
4670        // [{a: 2, b: null}]
4671        let values = list_builder.values();
4672        values
4673            .field_builder::<Int32Builder>(0)
4674            .unwrap()
4675            .append_value(2);
4676        values
4677            .field_builder::<Int32Builder>(1)
4678            .unwrap()
4679            .append_null();
4680        values.append(true);
4681        list_builder.append(true);
4682
4683        let array = Arc::new(list_builder.finish());
4684
4685        RoundTripTest::new(array).run();
4686    }
4687
4688    fn row_group_sizes(metadata: &ParquetMetaData) -> Vec<i64> {
4689        metadata.row_groups().iter().map(|x| x.num_rows()).collect()
4690    }
4691
4692    #[test]
4693    fn test_aggregates_records() {
4694        let arrays = [
4695            Int32Array::from((0..100).collect::<Vec<_>>()),
4696            Int32Array::from((0..50).collect::<Vec<_>>()),
4697            Int32Array::from((200..500).collect::<Vec<_>>()),
4698        ];
4699
4700        let schema = Arc::new(Schema::new(vec![Field::new(
4701            "int",
4702            ArrowDataType::Int32,
4703            false,
4704        )]));
4705
4706        let file = tempfile::tempfile().unwrap();
4707
4708        let props = WriterProperties::builder()
4709            .set_max_row_group_row_count(Some(200))
4710            .build();
4711
4712        let mut writer =
4713            ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
4714
4715        for array in arrays {
4716            let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
4717            writer.write(&batch).unwrap();
4718        }
4719
4720        writer.close().unwrap();
4721
4722        let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
4723        assert_eq!(&row_group_sizes(builder.metadata()), &[200, 200, 50]);
4724
4725        let batches = builder
4726            .with_batch_size(100)
4727            .build()
4728            .unwrap()
4729            .collect::<ArrowResult<Vec<_>>>()
4730            .unwrap();
4731
4732        assert_eq!(batches.len(), 5);
4733        assert!(batches.iter().all(|x| x.num_columns() == 1));
4734
4735        let batch_sizes: Vec<_> = batches.iter().map(|x| x.num_rows()).collect();
4736
4737        assert_eq!(&batch_sizes, &[100, 100, 100, 100, 50]);
4738
4739        let values: Vec<_> = batches
4740            .iter()
4741            .flat_map(|x| {
4742                x.column(0)
4743                    .as_any()
4744                    .downcast_ref::<Int32Array>()
4745                    .unwrap()
4746                    .values()
4747                    .iter()
4748                    .cloned()
4749            })
4750            .collect();
4751
4752        let expected_values: Vec<_> = [0..100, 0..50, 200..500].into_iter().flatten().collect();
4753        assert_eq!(&values, &expected_values)
4754    }
4755
4756    #[test]
4757    fn complex_aggregate() {
4758        // Tests aggregating nested data
4759        let field_a = Arc::new(Field::new("leaf_a", DataType::Int32, false));
4760        let field_b = Arc::new(Field::new("leaf_b", DataType::Int32, true));
4761        let struct_a = Arc::new(Field::new(
4762            "struct_a",
4763            DataType::Struct(vec![field_a.clone(), field_b.clone()].into()),
4764            true,
4765        ));
4766
4767        let list_a = Arc::new(Field::new("list", DataType::List(struct_a), true));
4768        let struct_b = Arc::new(Field::new(
4769            "struct_b",
4770            DataType::Struct(vec![list_a.clone()].into()),
4771            false,
4772        ));
4773
4774        let schema = Arc::new(Schema::new(vec![struct_b]));
4775
4776        // create nested data
4777        let field_a_array = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
4778        let field_b_array =
4779            Int32Array::from_iter(vec![Some(1), None, Some(2), None, None, Some(6)]);
4780
4781        let struct_a_array = StructArray::from(vec![
4782            (field_a.clone(), Arc::new(field_a_array) as ArrayRef),
4783            (field_b.clone(), Arc::new(field_b_array) as ArrayRef),
4784        ]);
4785
4786        let list_data = ArrayDataBuilder::new(list_a.data_type().clone())
4787            .len(5)
4788            .add_buffer(Buffer::from_iter(vec![
4789                0_i32, 1_i32, 1_i32, 3_i32, 3_i32, 5_i32,
4790            ]))
4791            .null_bit_buffer(Some(Buffer::from_iter(vec![
4792                true, false, true, false, true,
4793            ])))
4794            .child_data(vec![struct_a_array.into_data()])
4795            .build()
4796            .unwrap();
4797
4798        let list_a_array = Arc::new(ListArray::from(list_data)) as ArrayRef;
4799        let struct_b_array = StructArray::from(vec![(list_a.clone(), list_a_array)]);
4800
4801        let batch1 =
4802            RecordBatch::try_from_iter(vec![("struct_b", Arc::new(struct_b_array) as ArrayRef)])
4803                .unwrap();
4804
4805        let field_a_array = Int32Array::from(vec![6, 7, 8, 9, 10]);
4806        let field_b_array = Int32Array::from_iter(vec![None, None, None, Some(1), None]);
4807
4808        let struct_a_array = StructArray::from(vec![
4809            (field_a, Arc::new(field_a_array) as ArrayRef),
4810            (field_b, Arc::new(field_b_array) as ArrayRef),
4811        ]);
4812
4813        let list_data = ArrayDataBuilder::new(list_a.data_type().clone())
4814            .len(2)
4815            .add_buffer(Buffer::from_iter(vec![0_i32, 4_i32, 5_i32]))
4816            .child_data(vec![struct_a_array.into_data()])
4817            .build()
4818            .unwrap();
4819
4820        let list_a_array = Arc::new(ListArray::from(list_data)) as ArrayRef;
4821        let struct_b_array = StructArray::from(vec![(list_a, list_a_array)]);
4822
4823        let batch2 =
4824            RecordBatch::try_from_iter(vec![("struct_b", Arc::new(struct_b_array) as ArrayRef)])
4825                .unwrap();
4826
4827        let batches = &[batch1, batch2];
4828
4829        // Verify data is as expected
4830
4831        let expected = r"
4832            +-------------------------------------------------------------------------------------------------------+
4833            | struct_b                                                                                              |
4834            +-------------------------------------------------------------------------------------------------------+
4835            | {list: [{leaf_a: 1, leaf_b: 1}]}                                                                      |
4836            | {list: }                                                                                              |
4837            | {list: [{leaf_a: 2, leaf_b: }, {leaf_a: 3, leaf_b: 2}]}                                               |
4838            | {list: }                                                                                              |
4839            | {list: [{leaf_a: 4, leaf_b: }, {leaf_a: 5, leaf_b: }]}                                                |
4840            | {list: [{leaf_a: 6, leaf_b: }, {leaf_a: 7, leaf_b: }, {leaf_a: 8, leaf_b: }, {leaf_a: 9, leaf_b: 1}]} |
4841            | {list: [{leaf_a: 10, leaf_b: }]}                                                                      |
4842            +-------------------------------------------------------------------------------------------------------+
4843        ".trim().split('\n').map(|x| x.trim()).collect::<Vec<_>>().join("\n");
4844
4845        let actual = pretty_format_batches(batches).unwrap().to_string();
4846        assert_eq!(actual, expected);
4847
4848        // Write data
4849        let file = tempfile::tempfile().unwrap();
4850        let props = WriterProperties::builder()
4851            .set_max_row_group_row_count(Some(6))
4852            .build();
4853
4854        let mut writer =
4855            ArrowWriter::try_new(file.try_clone().unwrap(), schema, Some(props)).unwrap();
4856
4857        for batch in batches {
4858            writer.write(batch).unwrap();
4859        }
4860        writer.close().unwrap();
4861
4862        // Read Data
4863        // Should have written entire first batch and first row of second to the first row group
4864        // leaving a single row in the second row group
4865
4866        let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
4867        assert_eq!(&row_group_sizes(builder.metadata()), &[6, 1]);
4868
4869        let batches = builder
4870            .with_batch_size(2)
4871            .build()
4872            .unwrap()
4873            .collect::<ArrowResult<Vec<_>>>()
4874            .unwrap();
4875
4876        assert_eq!(batches.len(), 4);
4877        let batch_counts: Vec<_> = batches.iter().map(|x| x.num_rows()).collect();
4878        assert_eq!(&batch_counts, &[2, 2, 2, 1]);
4879
4880        let actual = pretty_format_batches(&batches).unwrap().to_string();
4881        assert_eq!(actual, expected);
4882    }
4883
4884    #[test]
4885    fn test_arrow_writer_metadata() {
4886        let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
4887        let file_schema = batch_schema.clone().with_metadata([("foo", "bar")]);
4888
4889        let batch = RecordBatch::try_new(
4890            Arc::new(batch_schema),
4891            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
4892        )
4893        .unwrap();
4894
4895        let mut buf = Vec::with_capacity(1024);
4896        let mut writer = ArrowWriter::try_new(&mut buf, Arc::new(file_schema), None).unwrap();
4897        writer.write(&batch).unwrap();
4898        writer.close().unwrap();
4899    }
4900
4901    #[test]
4902    fn test_arrow_writer_nullable() {
4903        let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
4904        let file_schema = Schema::new(vec![Field::new("int32", DataType::Int32, true)]);
4905        let file_schema = Arc::new(file_schema);
4906
4907        let batch = RecordBatch::try_new(
4908            Arc::new(batch_schema),
4909            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
4910        )
4911        .unwrap();
4912
4913        let mut buf = Vec::with_capacity(1024);
4914        let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), None).unwrap();
4915        writer.write(&batch).unwrap();
4916        writer.close().unwrap();
4917
4918        let mut read = ParquetRecordBatchReader::try_new(Bytes::from(buf), 1024).unwrap();
4919        let back = read.next().unwrap().unwrap();
4920        assert_eq!(back.schema(), file_schema);
4921        assert_ne!(back.schema(), batch.schema());
4922        assert_eq!(back.column(0).as_ref(), batch.column(0).as_ref());
4923    }
4924
4925    #[test]
4926    fn in_progress_accounting() {
4927        // define schema
4928        let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
4929
4930        // create some data
4931        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
4932
4933        // build a record batch
4934        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
4935
4936        let mut writer = ArrowWriter::try_new(vec![], batch.schema(), None).unwrap();
4937
4938        // starts empty
4939        assert_eq!(writer.in_progress_size(), 0);
4940        assert_eq!(writer.in_progress_rows(), 0);
4941        assert_eq!(writer.memory_size(), 0);
4942        assert_eq!(writer.bytes_written(), 4); // Initial header
4943        writer.write(&batch).unwrap();
4944
4945        // updated on write
4946        let initial_size = writer.in_progress_size();
4947        assert!(initial_size > 0);
4948        assert_eq!(writer.in_progress_rows(), 5);
4949        let initial_memory = writer.memory_size();
4950        assert!(initial_memory > 0);
4951        // memory estimate is larger than estimated encoded size
4952        assert!(
4953            initial_size <= initial_memory,
4954            "{initial_size} <= {initial_memory}"
4955        );
4956
4957        // updated on second write
4958        writer.write(&batch).unwrap();
4959        assert!(writer.in_progress_size() > initial_size);
4960        assert_eq!(writer.in_progress_rows(), 10);
4961        assert!(writer.memory_size() > initial_memory);
4962        assert!(
4963            writer.in_progress_size() <= writer.memory_size(),
4964            "in_progress_size {} <= memory_size {}",
4965            writer.in_progress_size(),
4966            writer.memory_size()
4967        );
4968
4969        // in progress tracking is cleared, but the overall data written is updated
4970        let pre_flush_bytes_written = writer.bytes_written();
4971        writer.flush().unwrap();
4972        assert_eq!(writer.in_progress_size(), 0);
4973        assert_eq!(writer.memory_size(), 0);
4974        assert!(writer.bytes_written() > pre_flush_bytes_written);
4975
4976        writer.close().unwrap();
4977    }
4978
4979    #[test]
4980    fn test_writer_all_null() {
4981        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
4982        let b = Int32Array::new(vec![0; 5].into(), Some(NullBuffer::new_null(5)));
4983        let batch = RecordBatch::try_from_iter(vec![
4984            ("a", Arc::new(a) as ArrayRef),
4985            ("b", Arc::new(b) as ArrayRef),
4986        ])
4987        .unwrap();
4988
4989        let mut buf = Vec::with_capacity(1024);
4990        let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), None).unwrap();
4991        writer.write(&batch).unwrap();
4992        writer.close().unwrap();
4993
4994        let bytes = Bytes::from(buf);
4995        let options = ReadOptionsBuilder::new().with_page_index().build();
4996        let reader = SerializedFileReader::new_with_options(bytes, options).unwrap();
4997        let index = reader.metadata().offset_index().unwrap();
4998
4999        assert_eq!(index.len(), 1);
5000        assert_eq!(index[0].len(), 2); // 2 columns
5001        assert_eq!(index[0][0].page_locations().len(), 1); // 1 page
5002        assert_eq!(index[0][1].page_locations().len(), 1); // 1 page
5003    }
5004
5005    #[test]
5006    fn test_disabled_statistics_with_page() {
5007        let file_schema = Schema::new(vec![
5008            Field::new("a", DataType::Utf8, true),
5009            Field::new("b", DataType::Utf8, true),
5010        ]);
5011        let file_schema = Arc::new(file_schema);
5012
5013        let batch = RecordBatch::try_new(
5014            file_schema.clone(),
5015            vec![
5016                Arc::new(StringArray::from(vec!["a", "b", "c", "d"])) as _,
5017                Arc::new(StringArray::from(vec!["w", "x", "y", "z"])) as _,
5018            ],
5019        )
5020        .unwrap();
5021
5022        let props = WriterProperties::builder()
5023            .set_statistics_enabled(EnabledStatistics::None)
5024            .set_column_statistics_enabled("a".into(), EnabledStatistics::Page)
5025            .build();
5026
5027        let mut buf = Vec::with_capacity(1024);
5028        let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), Some(props)).unwrap();
5029        writer.write(&batch).unwrap();
5030
5031        let metadata = writer.close().unwrap();
5032        assert_eq!(metadata.num_row_groups(), 1);
5033        let row_group = metadata.row_group(0);
5034        assert_eq!(row_group.num_columns(), 2);
5035        // Column "a" has both offset and column index, as requested
5036        assert!(row_group.column(0).offset_index_offset().is_some());
5037        assert!(row_group.column(0).column_index_offset().is_some());
5038        // Column "b" should only have offset index
5039        assert!(row_group.column(1).offset_index_offset().is_some());
5040        assert!(row_group.column(1).column_index_offset().is_none());
5041
5042        let options = ReadOptionsBuilder::new().with_page_index().build();
5043        let reader = SerializedFileReader::new_with_options(Bytes::from(buf), options).unwrap();
5044
5045        let row_group = reader.get_row_group(0).unwrap();
5046        let a_col = row_group.metadata().column(0);
5047        let b_col = row_group.metadata().column(1);
5048
5049        // Column chunk of column "a" should have chunk level statistics
5050        if let Statistics::ByteArray(byte_array_stats) = a_col.statistics().unwrap() {
5051            let min = byte_array_stats.min_opt().unwrap();
5052            let max = byte_array_stats.max_opt().unwrap();
5053
5054            assert_eq!(min.as_bytes(), b"a");
5055            assert_eq!(max.as_bytes(), b"d");
5056        } else {
5057            panic!("expecting Statistics::ByteArray");
5058        }
5059
5060        // The column chunk for column "b" shouldn't have statistics
5061        assert!(b_col.statistics().is_none());
5062
5063        let offset_index = reader.metadata().offset_index().unwrap();
5064        assert_eq!(offset_index.len(), 1); // 1 row group
5065        assert_eq!(offset_index[0].len(), 2); // 2 columns
5066
5067        let column_index = reader.metadata().column_index().unwrap();
5068        assert_eq!(column_index.len(), 1); // 1 row group
5069        assert_eq!(column_index[0].len(), 2); // 2 columns
5070
5071        let a_idx = &column_index[0][0];
5072        assert!(
5073            matches!(a_idx, ColumnIndexMetaData::BYTE_ARRAY(_)),
5074            "{a_idx:?}"
5075        );
5076        let b_idx = &column_index[0][1];
5077        assert!(matches!(b_idx, ColumnIndexMetaData::NONE), "{b_idx:?}");
5078    }
5079
5080    #[test]
5081    fn test_disabled_statistics_with_chunk() {
5082        let file_schema = Schema::new(vec![
5083            Field::new("a", DataType::Utf8, true),
5084            Field::new("b", DataType::Utf8, true),
5085        ]);
5086        let file_schema = Arc::new(file_schema);
5087
5088        let batch = RecordBatch::try_new(
5089            file_schema.clone(),
5090            vec![
5091                Arc::new(StringArray::from(vec!["a", "b", "c", "d"])) as _,
5092                Arc::new(StringArray::from(vec!["w", "x", "y", "z"])) as _,
5093            ],
5094        )
5095        .unwrap();
5096
5097        let props = WriterProperties::builder()
5098            .set_statistics_enabled(EnabledStatistics::None)
5099            .set_column_statistics_enabled("a".into(), EnabledStatistics::Chunk)
5100            .build();
5101
5102        let mut buf = Vec::with_capacity(1024);
5103        let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), Some(props)).unwrap();
5104        writer.write(&batch).unwrap();
5105
5106        let metadata = writer.close().unwrap();
5107        assert_eq!(metadata.num_row_groups(), 1);
5108        let row_group = metadata.row_group(0);
5109        assert_eq!(row_group.num_columns(), 2);
5110        // Column "a" should only have offset index
5111        assert!(row_group.column(0).offset_index_offset().is_some());
5112        assert!(row_group.column(0).column_index_offset().is_none());
5113        // Column "b" should only have offset index
5114        assert!(row_group.column(1).offset_index_offset().is_some());
5115        assert!(row_group.column(1).column_index_offset().is_none());
5116
5117        let options = ReadOptionsBuilder::new().with_page_index().build();
5118        let reader = SerializedFileReader::new_with_options(Bytes::from(buf), options).unwrap();
5119
5120        let row_group = reader.get_row_group(0).unwrap();
5121        let a_col = row_group.metadata().column(0);
5122        let b_col = row_group.metadata().column(1);
5123
5124        // Column chunk of column "a" should have chunk level statistics
5125        if let Statistics::ByteArray(byte_array_stats) = a_col.statistics().unwrap() {
5126            let min = byte_array_stats.min_opt().unwrap();
5127            let max = byte_array_stats.max_opt().unwrap();
5128
5129            assert_eq!(min.as_bytes(), b"a");
5130            assert_eq!(max.as_bytes(), b"d");
5131        } else {
5132            panic!("expecting Statistics::ByteArray");
5133        }
5134
5135        // The column chunk for column "b"  shouldn't have statistics
5136        assert!(b_col.statistics().is_none());
5137
5138        let column_index = reader.metadata().column_index().unwrap();
5139        assert_eq!(column_index.len(), 1); // 1 row group
5140        assert_eq!(column_index[0].len(), 2); // 2 columns
5141
5142        let a_idx = &column_index[0][0];
5143        assert!(matches!(a_idx, ColumnIndexMetaData::NONE), "{a_idx:?}");
5144        let b_idx = &column_index[0][1];
5145        assert!(matches!(b_idx, ColumnIndexMetaData::NONE), "{b_idx:?}");
5146    }
5147
5148    #[test]
5149    fn test_arrow_writer_skip_metadata() {
5150        let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
5151        let file_schema = Arc::new(batch_schema.clone());
5152
5153        let batch = RecordBatch::try_new(
5154            Arc::new(batch_schema),
5155            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5156        )
5157        .unwrap();
5158        let skip_options = ArrowWriterOptions::new().with_skip_arrow_metadata(true);
5159
5160        let mut buf = Vec::with_capacity(1024);
5161        let mut writer =
5162            ArrowWriter::try_new_with_options(&mut buf, file_schema.clone(), skip_options).unwrap();
5163        writer.write(&batch).unwrap();
5164        writer.close().unwrap();
5165
5166        let bytes = Bytes::from(buf);
5167        let reader_builder = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap();
5168        assert_eq!(file_schema, *reader_builder.schema());
5169        if let Some(key_value_metadata) = reader_builder
5170            .metadata()
5171            .file_metadata()
5172            .key_value_metadata()
5173        {
5174            assert!(
5175                !key_value_metadata
5176                    .iter()
5177                    .any(|kv| kv.key.as_str() == ARROW_SCHEMA_META_KEY)
5178            );
5179        }
5180    }
5181
5182    #[test]
5183    fn test_arrow_writer_skip_path_in_schema() {
5184        let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
5185        let file_schema = Arc::new(batch_schema.clone());
5186
5187        let batch = RecordBatch::try_new(
5188            Arc::new(batch_schema),
5189            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5190        )
5191        .unwrap();
5192
5193        // default options should still write path_in_schema
5194        let skip_options = ArrowWriterOptions::new();
5195
5196        let mut buf = Vec::with_capacity(1024);
5197        let mut writer =
5198            ArrowWriter::try_new_with_options(&mut buf, file_schema.clone(), skip_options).unwrap();
5199        writer.write(&batch).unwrap();
5200        writer.close().unwrap();
5201
5202        // override to not write path_in_schema
5203        let skip_options = ArrowWriterOptions::new().with_properties(
5204            WriterProperties::builder()
5205                .set_write_path_in_schema(false)
5206                .build(),
5207        );
5208
5209        let mut buf2 = Vec::with_capacity(1024);
5210        let mut writer =
5211            ArrowWriter::try_new_with_options(&mut buf2, file_schema.clone(), skip_options)
5212                .unwrap();
5213        writer.write(&batch).unwrap();
5214        writer.close().unwrap();
5215
5216        // buf2 should be a bit smaller due to lack of path_in_schema
5217        assert!(buf.len() > buf2.len());
5218    }
5219
5220    #[test]
5221    fn mismatched_schemas() {
5222        let batch_schema = Schema::new(vec![Field::new("count", DataType::Int32, false)]);
5223        let file_schema = Arc::new(Schema::new(vec![Field::new(
5224            "temperature",
5225            DataType::Float64,
5226            false,
5227        )]));
5228
5229        let batch = RecordBatch::try_new(
5230            Arc::new(batch_schema),
5231            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5232        )
5233        .unwrap();
5234
5235        let mut buf = Vec::with_capacity(1024);
5236        let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), None).unwrap();
5237
5238        let err = writer.write(&batch).unwrap_err().to_string();
5239        assert_eq!(
5240            err,
5241            "Arrow: Incompatible type. Field 'temperature' has type Float64, array has type Int32"
5242        );
5243    }
5244
5245    #[test]
5246    // https://github.com/apache/arrow-rs/issues/6988
5247    fn test_roundtrip_empty_schema() {
5248        // create empty record batch with empty schema
5249        let empty_batch = RecordBatch::try_new_with_options(
5250            Arc::new(Schema::empty()),
5251            vec![],
5252            &RecordBatchOptions::default().with_row_count(Some(0)),
5253        )
5254        .unwrap();
5255
5256        // write to parquet
5257        let mut parquet_bytes: Vec<u8> = Vec::new();
5258        let mut writer =
5259            ArrowWriter::try_new(&mut parquet_bytes, empty_batch.schema(), None).unwrap();
5260        writer.write(&empty_batch).unwrap();
5261        writer.close().unwrap();
5262
5263        // read from parquet
5264        let bytes = Bytes::from(parquet_bytes);
5265        let reader = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap();
5266        assert_eq!(reader.schema(), &empty_batch.schema());
5267        let batches: Vec<_> = reader
5268            .build()
5269            .unwrap()
5270            .collect::<ArrowResult<Vec<_>>>()
5271            .unwrap();
5272        assert_eq!(batches.len(), 0);
5273    }
5274
5275    #[test]
5276    fn test_page_stats_not_written_by_default() {
5277        let string_field = Field::new("a", DataType::Utf8, false);
5278        let schema = Schema::new(vec![string_field]);
5279        let raw_string_values = vec!["Blart Versenwald III"];
5280        let string_values = StringArray::from(raw_string_values.clone());
5281        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(string_values)]).unwrap();
5282
5283        let props = WriterProperties::builder()
5284            .set_statistics_enabled(EnabledStatistics::Page)
5285            .set_dictionary_enabled(false)
5286            .set_encoding(Encoding::PLAIN)
5287            .set_compression(crate::basic::Compression::UNCOMPRESSED)
5288            .build();
5289
5290        let file = roundtrip_opts(&batch, props);
5291
5292        // read file and decode page headers
5293        // Note: use the thrift API as there is no Rust API to access the statistics in the page headers
5294
5295        // decode first page header
5296        let first_page = &file[4..];
5297        let mut prot = ThriftSliceInputProtocol::new(first_page);
5298        let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5299        let stats = hdr.data_page_header.unwrap().statistics;
5300
5301        assert!(stats.is_none());
5302    }
5303
5304    #[test]
5305    fn test_page_stats_when_enabled() {
5306        let string_field = Field::new("a", DataType::Utf8, false);
5307        let schema = Schema::new(vec![string_field]);
5308        let raw_string_values = vec!["Blart Versenwald III", "Andrew Lamb"];
5309        let string_values = StringArray::from(raw_string_values.clone());
5310        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(string_values)]).unwrap();
5311
5312        let props = WriterProperties::builder()
5313            .set_statistics_enabled(EnabledStatistics::Page)
5314            .set_dictionary_enabled(false)
5315            .set_encoding(Encoding::PLAIN)
5316            .set_write_page_header_statistics(true)
5317            .set_compression(crate::basic::Compression::UNCOMPRESSED)
5318            .build();
5319
5320        let file = roundtrip_opts(&batch, props);
5321
5322        // read file and decode page headers
5323        // Note: use the thrift API as there is no Rust API to access the statistics in the page headers
5324
5325        // decode first page header
5326        let first_page = &file[4..];
5327        let mut prot = ThriftSliceInputProtocol::new(first_page);
5328        let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5329        let stats = hdr.data_page_header.unwrap().statistics;
5330
5331        let stats = stats.unwrap();
5332        // check that min/max were actually written to the page
5333        assert!(stats.is_max_value_exact.unwrap());
5334        assert!(stats.is_min_value_exact.unwrap());
5335        assert_eq!(stats.max_value.unwrap(), "Blart Versenwald III".as_bytes());
5336        assert_eq!(stats.min_value.unwrap(), "Andrew Lamb".as_bytes());
5337    }
5338
5339    #[test]
5340    fn test_page_stats_truncation() {
5341        let string_field = Field::new("a", DataType::Utf8, false);
5342        let binary_field = Field::new("b", DataType::Binary, false);
5343        let schema = Schema::new(vec![string_field, binary_field]);
5344
5345        let raw_string_values = vec!["Blart Versenwald III"];
5346        let raw_binary_values = [b"Blart Versenwald III".to_vec()];
5347        let raw_binary_value_refs = raw_binary_values
5348            .iter()
5349            .map(|x| x.as_slice())
5350            .collect::<Vec<_>>();
5351
5352        let string_values = StringArray::from(raw_string_values.clone());
5353        let binary_values = BinaryArray::from(raw_binary_value_refs);
5354        let batch = RecordBatch::try_new(
5355            Arc::new(schema),
5356            vec![Arc::new(string_values), Arc::new(binary_values)],
5357        )
5358        .unwrap();
5359
5360        let props = WriterProperties::builder()
5361            .set_statistics_truncate_length(Some(2))
5362            .set_dictionary_enabled(false)
5363            .set_encoding(Encoding::PLAIN)
5364            .set_write_page_header_statistics(true)
5365            .set_compression(crate::basic::Compression::UNCOMPRESSED)
5366            .build();
5367
5368        let file = roundtrip_opts(&batch, props);
5369
5370        // read file and decode page headers
5371        // Note: use the thrift API as there is no Rust API to access the statistics in the page headers
5372
5373        // decode first page header
5374        let first_page = &file[4..];
5375        let mut prot = ThriftSliceInputProtocol::new(first_page);
5376        let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5377        let stats = hdr.data_page_header.unwrap().statistics;
5378        assert!(stats.is_some());
5379        let stats = stats.unwrap();
5380        // check that min/max were properly truncated
5381        assert!(!stats.is_max_value_exact.unwrap());
5382        assert!(!stats.is_min_value_exact.unwrap());
5383        assert_eq!(stats.max_value.unwrap(), "Bm".as_bytes());
5384        assert_eq!(stats.min_value.unwrap(), "Bl".as_bytes());
5385
5386        // check second page now
5387        let second_page = &prot.as_slice()[hdr.compressed_page_size as usize..];
5388        let mut prot = ThriftSliceInputProtocol::new(second_page);
5389        let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5390        let stats = hdr.data_page_header.unwrap().statistics;
5391        assert!(stats.is_some());
5392        let stats = stats.unwrap();
5393        // check that min/max were properly truncated
5394        assert!(!stats.is_max_value_exact.unwrap());
5395        assert!(!stats.is_min_value_exact.unwrap());
5396        assert_eq!(stats.max_value.unwrap(), "Bm".as_bytes());
5397        assert_eq!(stats.min_value.unwrap(), "Bl".as_bytes());
5398    }
5399
5400    #[test]
5401    fn test_page_encoding_statistics_roundtrip() {
5402        let batch_schema = Schema::new(vec![Field::new(
5403            "int32",
5404            arrow_schema::DataType::Int32,
5405            false,
5406        )]);
5407
5408        let batch = RecordBatch::try_new(
5409            Arc::new(batch_schema.clone()),
5410            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5411        )
5412        .unwrap();
5413
5414        let mut file: File = tempfile::tempfile().unwrap();
5415        let mut writer = ArrowWriter::try_new(&mut file, Arc::new(batch_schema), None).unwrap();
5416        writer.write(&batch).unwrap();
5417        let file_metadata = writer.close().unwrap();
5418
5419        assert_eq!(file_metadata.num_row_groups(), 1);
5420        assert_eq!(file_metadata.row_group(0).num_columns(), 1);
5421        assert!(
5422            file_metadata
5423                .row_group(0)
5424                .column(0)
5425                .page_encoding_stats()
5426                .is_some()
5427        );
5428        let chunk_page_stats = file_metadata
5429            .row_group(0)
5430            .column(0)
5431            .page_encoding_stats()
5432            .unwrap();
5433
5434        // check that the read metadata is also correct
5435        let options = ReadOptionsBuilder::new()
5436            .with_page_index()
5437            .with_encoding_stats_as_mask(false)
5438            .build();
5439        let reader = SerializedFileReader::new_with_options(file, options).unwrap();
5440
5441        let rowgroup = reader.get_row_group(0).expect("row group missing");
5442        assert_eq!(rowgroup.num_columns(), 1);
5443        let column = rowgroup.metadata().column(0);
5444        assert!(column.page_encoding_stats().is_some());
5445        let file_page_stats = column.page_encoding_stats().unwrap();
5446        assert_eq!(chunk_page_stats, file_page_stats);
5447    }
5448
5449    #[test]
5450    fn test_different_dict_page_size_limit() {
5451        let array = Arc::new(Int64Array::from_iter(0..1024 * 1024));
5452        let schema = Arc::new(Schema::new(vec![
5453            Field::new("col0", arrow_schema::DataType::Int64, false),
5454            Field::new("col1", arrow_schema::DataType::Int64, false),
5455        ]));
5456        let batch =
5457            arrow_array::RecordBatch::try_new(schema.clone(), vec![array.clone(), array]).unwrap();
5458
5459        let props = WriterProperties::builder()
5460            .set_dictionary_page_size_limit(1024 * 1024)
5461            .set_column_dictionary_page_size_limit(ColumnPath::from("col1"), 1024 * 1024 * 4)
5462            .build();
5463        let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5464        writer.write(&batch).unwrap();
5465        let data = Bytes::from(writer.into_inner().unwrap());
5466
5467        let mut metadata = ParquetMetaDataReader::new();
5468        metadata.try_parse(&data).unwrap();
5469        let metadata = metadata.finish().unwrap();
5470        let col0_meta = metadata.row_group(0).column(0);
5471        let col1_meta = metadata.row_group(0).column(1);
5472
5473        let get_dict_page_size = move |meta: &ColumnChunkMetaData| {
5474            let mut reader =
5475                SerializedPageReader::new(Arc::new(data.clone()), meta, 0, None).unwrap();
5476            let page = reader.get_next_page().unwrap().unwrap();
5477            match page {
5478                Page::DictionaryPage { buf, .. } => buf.len(),
5479                _ => panic!("expected DictionaryPage"),
5480            }
5481        };
5482
5483        assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
5484        assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
5485    }
5486
5487    #[test]
5488    fn test_arrow_writer_granular_mode_roundtrip() {
5489        // Granular mode subdivides chunks and writes more pages than the
5490        // default batched path. Make sure the data we write back is
5491        // bit-identical to what went in — page-count assertions elsewhere
5492        // only prove pages were cut, not that the encoded data is correct.
5493        //
5494        // Mix value sizes so that the cumulative-byte-budget cutoff
5495        // lands mid-chunk, exercising both batched and granular paths
5496        // within the same `write_batch_internal` call.
5497        let small = "tiny".to_string();
5498        let big = "x".repeat(64 * 1024);
5499        let strings: Vec<String> = (0..256)
5500            .map(|i| {
5501                if i % 16 == 0 {
5502                    big.clone()
5503                } else {
5504                    small.clone()
5505                }
5506            })
5507            .collect();
5508
5509        let schema = Arc::new(Schema::new(vec![Field::new(
5510            "col",
5511            ArrowDataType::Utf8,
5512            false,
5513        )]));
5514        let batch = RecordBatch::try_new(
5515            schema.clone(),
5516            vec![Arc::new(StringArray::from(strings.clone())) as _],
5517        )
5518        .unwrap();
5519
5520        let props = WriterProperties::builder()
5521            .set_dictionary_enabled(false)
5522            .set_data_page_size_limit(16 * 1024)
5523            .build();
5524        let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5525        writer.write(&batch).unwrap();
5526        let data = Bytes::from(writer.into_inner().unwrap());
5527
5528        let mut reader = ParquetRecordBatchReader::try_new(data, 1024).unwrap();
5529        let read = reader.next().unwrap().unwrap();
5530        assert!(reader.next().is_none(), "expected one batch");
5531        let col = read
5532            .column(0)
5533            .as_any()
5534            .downcast_ref::<StringArray>()
5535            .unwrap();
5536        assert_eq!(col.len(), strings.len());
5537        for (i, expected) in strings.iter().enumerate() {
5538            assert_eq!(
5539                col.value(i),
5540                expected.as_str(),
5541                "value mismatch at index {i}"
5542            );
5543        }
5544    }
5545
5546    #[test]
5547    fn test_arrow_writer_all_null_string_column() {
5548        // The `LevelDataRef::value_count` Uniform branch with
5549        // `value != max_def` (entirely-null chunk) must return 0 so the
5550        // sub-batch sizer short-circuits to batch mode without trying
5551        // to estimate byte budgets for non-existent values.
5552        let num_rows = 1024;
5553        let schema = Arc::new(Schema::new(vec![Field::new(
5554            "col",
5555            ArrowDataType::Utf8,
5556            true,
5557        )]));
5558        let nulls: Vec<Option<&str>> = vec![None; num_rows];
5559        let batch = RecordBatch::try_new(
5560            schema.clone(),
5561            vec![Arc::new(StringArray::from(nulls)) as _],
5562        )
5563        .unwrap();
5564
5565        let props = WriterProperties::builder()
5566            .set_dictionary_enabled(false)
5567            .set_data_page_size_limit(16 * 1024)
5568            .build();
5569        let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5570        writer.write(&batch).unwrap();
5571        let data = Bytes::from(writer.into_inner().unwrap());
5572
5573        // Re-parse the file: row group has one column, every row is
5574        // null, all data pages report `num_rows / page_count` rows.
5575        let mut metadata = ParquetMetaDataReader::new();
5576        metadata.try_parse(&data).unwrap();
5577        let metadata = metadata.finish().unwrap();
5578        let row_group = metadata.row_group(0);
5579        let col_meta = row_group.column(0);
5580        assert_eq!(row_group.num_rows() as usize, num_rows);
5581        // Statistics record `null_count = num_rows` — proves every value
5582        // was written as null.
5583        if let Some(stats) = col_meta.statistics() {
5584            assert_eq!(
5585                stats.null_count_opt().unwrap_or(0) as usize,
5586                num_rows,
5587                "expected all-null column to report null_count = num_rows"
5588            );
5589        }
5590
5591        let mut reader =
5592            SerializedPageReader::new(Arc::new(data.clone()), col_meta, num_rows, None).unwrap();
5593        let mut total_values = 0u32;
5594        while let Some(page) = reader.get_next_page().unwrap() {
5595            if matches!(page, Page::DataPage { .. } | Page::DataPageV2 { .. }) {
5596                total_values += page.num_values();
5597            }
5598        }
5599        assert_eq!(
5600            total_values as usize, num_rows,
5601            "expected every level position to be represented in some page"
5602        );
5603    }
5604
5605    struct WriteBatchesShape {
5606        num_batches: usize,
5607        rows_per_batch: usize,
5608        row_size: usize,
5609    }
5610
5611    /// Helper function to write batches with the provided `WriteBatchesShape` into an `ArrowWriter`
5612    fn write_batches(
5613        WriteBatchesShape {
5614            num_batches,
5615            rows_per_batch,
5616            row_size,
5617        }: WriteBatchesShape,
5618        props: WriterProperties,
5619    ) -> ParquetRecordBatchReaderBuilder<File> {
5620        let schema = Arc::new(Schema::new(vec![Field::new(
5621            "str",
5622            ArrowDataType::Utf8,
5623            false,
5624        )]));
5625        let file = tempfile::tempfile().unwrap();
5626        let mut writer =
5627            ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
5628
5629        for batch_idx in 0..num_batches {
5630            let strings: Vec<String> = (0..rows_per_batch)
5631                .map(|i| format!("{:0>width$}", batch_idx * 10 + i, width = row_size))
5632                .collect();
5633            let array = StringArray::from(strings);
5634            let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
5635            writer.write(&batch).unwrap();
5636        }
5637        writer.close().unwrap();
5638        ParquetRecordBatchReaderBuilder::try_new(file).unwrap()
5639    }
5640
5641    #[test]
5642    // When both limits are None, all data should go into a single row group
5643    fn test_row_group_limit_none_writes_single_row_group() {
5644        let props = WriterProperties::builder()
5645            .set_max_row_group_row_count(None)
5646            .set_max_row_group_bytes(None)
5647            .build();
5648
5649        let builder = write_batches(
5650            WriteBatchesShape {
5651                num_batches: 1,
5652                rows_per_batch: 1000,
5653                row_size: 4,
5654            },
5655            props,
5656        );
5657
5658        assert_eq!(
5659            &row_group_sizes(builder.metadata()),
5660            &[1000],
5661            "With no limits, all rows should be in a single row group"
5662        );
5663    }
5664
5665    #[test]
5666    // When only max_row_group_size is set, respect the row limit
5667    fn test_row_group_limit_rows_only() {
5668        let props = WriterProperties::builder()
5669            .set_max_row_group_row_count(Some(300))
5670            .set_max_row_group_bytes(None)
5671            .build();
5672
5673        let builder = write_batches(
5674            WriteBatchesShape {
5675                num_batches: 1,
5676                rows_per_batch: 1000,
5677                row_size: 4,
5678            },
5679            props,
5680        );
5681
5682        assert_eq!(
5683            &row_group_sizes(builder.metadata()),
5684            &[300, 300, 300, 100],
5685            "Row groups should be split by row count"
5686        );
5687    }
5688
5689    #[test]
5690    // When only max_row_group_bytes is set, respect the byte limit
5691    fn test_row_group_limit_bytes_only() {
5692        let props = WriterProperties::builder()
5693            .set_max_row_group_row_count(None)
5694            // Set byte limit to approximately fit ~30 rows worth of data (~100 bytes each)
5695            .set_max_row_group_bytes(Some(3500))
5696            .build();
5697
5698        let builder = write_batches(
5699            WriteBatchesShape {
5700                num_batches: 10,
5701                rows_per_batch: 10,
5702                row_size: 100,
5703            },
5704            props,
5705        );
5706
5707        let sizes = row_group_sizes(builder.metadata());
5708
5709        assert!(
5710            sizes.len() > 1,
5711            "Should have multiple row groups due to byte limit, got {sizes:?}",
5712        );
5713
5714        let total_rows: i64 = sizes.iter().sum();
5715        assert_eq!(total_rows, 100, "Total rows should be preserved");
5716    }
5717
5718    #[test]
5719    // If an in-progress row group is already oversized, it should be flushed before writing more.
5720    fn test_row_group_limit_bytes_flushes_when_current_group_already_too_large() {
5721        let schema = Arc::new(Schema::new(vec![Field::new(
5722            "str",
5723            ArrowDataType::Utf8,
5724            false,
5725        )]));
5726        let file = tempfile::tempfile().unwrap();
5727
5728        // Start with no byte limit so we can intentionally build an oversized in-progress row group.
5729        let props = WriterProperties::builder()
5730            .set_max_row_group_row_count(None)
5731            .set_max_row_group_bytes(None)
5732            .build();
5733        let mut writer =
5734            ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
5735
5736        let first_array = StringArray::from(
5737            (0..10)
5738                .map(|i| format!("{:0>100}", i))
5739                .collect::<Vec<String>>(),
5740        );
5741        let first_batch =
5742            RecordBatch::try_new(schema.clone(), vec![Arc::new(first_array)]).unwrap();
5743        writer.write(&first_batch).unwrap();
5744        assert_eq!(writer.in_progress_rows(), 10);
5745
5746        // Tighten the limit below the current in-progress bytes to exercise:
5747        // `if current_bytes >= max_bytes { self.flush()?; ... }`
5748        writer.max_row_group_bytes = Some(1);
5749
5750        let second_array = StringArray::from(vec!["x".to_string()]);
5751        let second_batch =
5752            RecordBatch::try_new(schema.clone(), vec![Arc::new(second_array)]).unwrap();
5753        writer.write(&second_batch).unwrap();
5754        writer.close().unwrap();
5755        let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
5756
5757        assert_eq!(
5758            &row_group_sizes(builder.metadata()),
5759            &[10, 1],
5760            "The second write should flush an oversized in-progress row group first",
5761        );
5762    }
5763
5764    #[test]
5765    // When both limits are set, the row limit triggers first
5766    fn test_row_group_limit_both_row_wins_single_batch() {
5767        let props = WriterProperties::builder()
5768            .set_max_row_group_row_count(Some(200)) // Will trigger at 200 rows
5769            .set_max_row_group_bytes(Some(1024 * 1024)) // 1MB - won't trigger for small int data
5770            .build();
5771
5772        let builder = write_batches(
5773            WriteBatchesShape {
5774                num_batches: 1,
5775                row_size: 4,
5776                rows_per_batch: 1000,
5777            },
5778            props,
5779        );
5780
5781        assert_eq!(
5782            &row_group_sizes(builder.metadata()),
5783            &[200, 200, 200, 200, 200],
5784            "Row limit should trigger before byte limit"
5785        );
5786    }
5787
5788    #[test]
5789    // When both limits are set, the row limit triggers first
5790    fn test_row_group_limit_both_row_wins_multiple_batches() {
5791        let props = WriterProperties::builder()
5792            .set_max_row_group_row_count(Some(5)) // Will trigger every 5 rows
5793            .set_max_row_group_bytes(Some(9999)) // Won't trigger
5794            .build();
5795
5796        let builder = write_batches(
5797            WriteBatchesShape {
5798                num_batches: 10,
5799                rows_per_batch: 10,
5800                row_size: 100,
5801            },
5802            props,
5803        );
5804
5805        assert_eq!(
5806            &row_group_sizes(builder.metadata()),
5807            &[5; 20],
5808            "Row limit should trigger before byte limit"
5809        );
5810    }
5811
5812    #[test]
5813    // When both limits are set, the byte limit triggers first
5814    fn test_row_group_limit_both_bytes_wins() {
5815        let props = WriterProperties::builder()
5816            .set_max_row_group_row_count(Some(1000)) // Won't trigger for 100 rows
5817            .set_max_row_group_bytes(Some(3500)) // Will trigger at ~30-35 rows
5818            .build();
5819
5820        let builder = write_batches(
5821            WriteBatchesShape {
5822                num_batches: 10,
5823                rows_per_batch: 10,
5824                row_size: 100,
5825            },
5826            props,
5827        );
5828
5829        let sizes = row_group_sizes(builder.metadata());
5830
5831        assert!(
5832            sizes.len() > 1,
5833            "Byte limit should trigger before row limit, got {sizes:?}",
5834        );
5835
5836        assert!(
5837            sizes.iter().all(|&s| s < 1000),
5838            "No row group should hit the row limit"
5839        );
5840
5841        let total_rows: i64 = sizes.iter().sum();
5842        assert_eq!(total_rows, 100, "Total rows should be preserved");
5843    }
5844
5845    #[test]
5846    fn arrow_column_chunk_close_mut_drops_column_index() {
5847        use crate::arrow::ArrowSchemaConverter;
5848        use crate::file::writer::SerializedFileWriter;
5849
5850        let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)]));
5851        let props = Arc::new(
5852            WriterProperties::builder()
5853                .set_statistics_enabled(EnabledStatistics::Page)
5854                .build(),
5855        );
5856        let parquet_schema = ArrowSchemaConverter::new()
5857            .with_coerce_types(props.coerce_types())
5858            .convert(&schema)
5859            .unwrap();
5860
5861        let mut buf = Vec::with_capacity(1024);
5862        let mut writer =
5863            SerializedFileWriter::new(&mut buf, parquet_schema.root_schema_ptr(), props.clone())
5864                .unwrap();
5865
5866        let factory = ArrowRowGroupWriterFactory::new(&writer, Arc::clone(&schema));
5867        let mut col_writers = factory.create_column_writers(0).unwrap();
5868        let arr: ArrayRef = Arc::new(Int32Array::from_iter_values(0..64));
5869        for leaves in compute_leaves(schema.field(0), &arr).unwrap() {
5870            col_writers[0].write(&leaves).unwrap();
5871        }
5872        let mut chunk = col_writers.pop().unwrap().close().unwrap();
5873
5874        // Immutable accessor exposes the close result produced at close time.
5875        assert!(
5876            chunk.close().column_index.is_some(),
5877            "EnabledStatistics::Page should produce a column_index"
5878        );
5879
5880        // Mutable accessor lets callers drop the page-level index before append.
5881        chunk.close_mut().column_index = None;
5882        assert!(chunk.close().column_index.is_none());
5883
5884        let mut rg = writer.next_row_group().unwrap();
5885        chunk.append_to_row_group(&mut rg).unwrap();
5886        rg.close().unwrap();
5887        let file_meta = writer.close().unwrap();
5888
5889        // After dropping column_index, the resulting file records no column
5890        // index offset/length for this chunk.
5891        let cc = file_meta.row_group(0).column(0);
5892        assert!(cc.column_index_range().is_none());
5893    }
5894
5895    /// Writes a single-column RecordBatch to an in-memory Parquet buffer.
5896    fn write_column_to_bytes(array: ArrayRef) -> Bytes {
5897        let schema = Arc::new(Schema::new(vec![Field::new(
5898            "col",
5899            array.data_type().clone(),
5900            true,
5901        )]));
5902        let buf = get_bytes_after_close(
5903            schema.clone(),
5904            &RecordBatch::try_new(schema, vec![array]).unwrap(),
5905        );
5906        Bytes::from(buf)
5907    }
5908
5909    /// Reads column 0 from a single-row-group Parquet buffer, projecting it with the given schema.
5910    /// Passing a flat schema when the buffer was written from a REE array lets callers decode
5911    /// the physical values without the run-end encoding wrapper.
5912    fn read_column_with_schema(bytes: Bytes, schema: SchemaRef) -> ArrayRef {
5913        let opts = crate::arrow::arrow_reader::ArrowReaderOptions::new().with_schema(schema);
5914        ParquetRecordBatchReaderBuilder::try_new_with_options(bytes, opts)
5915            .unwrap()
5916            .build()
5917            .unwrap()
5918            .next()
5919            .unwrap()
5920            .unwrap()
5921            .column(0)
5922            .clone()
5923    }
5924
5925    fn ree_write_read_roundtrip(ree: ArrayRef, flat: ArrayRef) {
5926        let flat_schema = Arc::new(Schema::new(vec![Field::new(
5927            "col",
5928            flat.data_type().clone(),
5929            true,
5930        )]));
5931        let ree_bytes = write_column_to_bytes(ree);
5932        let flat_bytes = write_column_to_bytes(flat.clone());
5933        assert_eq!(
5934            ree_bytes, flat_bytes,
5935            "REE and flat bytes should be identical"
5936        );
5937
5938        let decoded_ree = read_column_with_schema(ree_bytes, flat_schema.clone());
5939        let decoded_flat = read_column_with_schema(flat_bytes, flat_schema);
5940
5941        assert_eq!(decoded_ree.as_ref(), flat.as_ref());
5942        assert_eq!(decoded_ree.as_ref(), decoded_flat.as_ref());
5943    }
5944
5945    #[test]
5946    fn ree_string() {
5947        let ree: ArrayRef = Arc::new(
5948            [Some("a"), Some("a"), None, Some("b"), Some("b")]
5949                .into_iter()
5950                .collect::<Int32RunArray>(),
5951        );
5952        let flat: ArrayRef = Arc::new(StringArray::from(vec![
5953            Some("a"),
5954            Some("a"),
5955            None,
5956            Some("b"),
5957            Some("b"),
5958        ]));
5959        ree_write_read_roundtrip(ree, flat);
5960    }
5961
5962    #[test]
5963    fn ree_int32() {
5964        let mut b = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
5965        for v in [Some(1), Some(1), None, Some(2), Some(2)] {
5966            b.append_option(v);
5967        }
5968        let ree: ArrayRef = Arc::new(b.finish());
5969        let flat: ArrayRef = Arc::new(Int32Array::from(vec![
5970            Some(1),
5971            Some(1),
5972            None,
5973            Some(2),
5974            Some(2),
5975        ]));
5976        ree_write_read_roundtrip(ree, flat);
5977    }
5978
5979    #[test]
5980    fn ree_bool() {
5981        // run_ends [3, 5, 7] → [T,T,T, null,null, F,F]
5982        let ree: ArrayRef = Arc::new(
5983            RunArray::try_new(
5984                &Int32Array::from(vec![3, 5, 7]),
5985                &BooleanArray::from(vec![Some(true), None, Some(false)]),
5986            )
5987            .unwrap(),
5988        );
5989        let flat: ArrayRef = Arc::new(BooleanArray::from(vec![
5990            Some(true),
5991            Some(true),
5992            Some(true),
5993            None,
5994            None,
5995            Some(false),
5996            Some(false),
5997        ]));
5998        ree_write_read_roundtrip(ree, flat);
5999    }
6000
6001    #[test]
6002    fn ree_fixed_size_binary() {
6003        let mk = |vals: &[Option<&[u8]>]| -> FixedSizeBinaryArray {
6004            let mut b = FixedSizeBinaryBuilder::new(2);
6005            for v in vals {
6006                match v {
6007                    Some(x) => b.append_value(x).unwrap(),
6008                    None => b.append_null(),
6009                }
6010            }
6011            b.finish()
6012        };
6013        // run_ends [2, 4, 6] → [aa,aa, null,null, bb,bb]
6014        let ree: ArrayRef = Arc::new(
6015            RunArray::try_new(
6016                &Int32Array::from(vec![2, 4, 6]),
6017                &mk(&[Some(b"aa"), None, Some(b"bb")]),
6018            )
6019            .unwrap(),
6020        );
6021        let flat: ArrayRef = Arc::new(mk(&[
6022            Some(b"aa"),
6023            Some(b"aa"),
6024            None,
6025            None,
6026            Some(b"bb"),
6027            Some(b"bb"),
6028        ]));
6029        ree_write_read_roundtrip(ree, flat);
6030    }
6031
6032    #[test]
6033    fn ree_single_run() {
6034        let ree: ArrayRef = Arc::new(["x", "x", "x"].into_iter().collect::<Int32RunArray>());
6035        let flat: ArrayRef = Arc::new(StringArray::from(vec!["x", "x", "x"]));
6036        ree_write_read_roundtrip(ree, flat);
6037    }
6038
6039    #[test]
6040    fn ree_float32() {
6041        // run_ends [2, 4, 5] → [1.0, 1.0, null, null, 2.5]
6042        let ree: ArrayRef = Arc::new(
6043            RunArray::try_new(
6044                &Int32Array::from(vec![2, 4, 5]),
6045                &Float32Array::from(vec![Some(1.0_f32), None, Some(2.5_f32)]),
6046            )
6047            .unwrap(),
6048        );
6049        let flat: ArrayRef = Arc::new(Float32Array::from(vec![
6050            Some(1.0_f32),
6051            Some(1.0_f32),
6052            None,
6053            None,
6054            Some(2.5_f32),
6055        ]));
6056        ree_write_read_roundtrip(ree, flat);
6057    }
6058
6059    #[test]
6060    fn ree_sliced() {
6061        // A sliced (non-zero offset) REE array: verify that get_physical_index
6062        // correctly accounts for the logical offset when expanding.
6063        // Full array: run_ends [3, 5, 7] → [a,a,a, b,b, c,c]
6064        // After slice(2, 5) the logical view is [a, b, b, c, c].
6065        let full: ArrayRef = Arc::new(
6066            RunArray::try_new(
6067                &Int32Array::from(vec![3, 5, 7]),
6068                &StringArray::from(vec!["a", "b", "c"]),
6069            )
6070            .unwrap(),
6071        );
6072        let sliced = full.slice(2, 5);
6073        let flat: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "b", "c", "c"]));
6074        ree_write_read_roundtrip(sliced, flat);
6075    }
6076
6077    #[test]
6078    fn ree_struct_with_ree_child() {
6079        // Struct with a REE string field and a REE int field — confirms
6080        // recursion visits every child and each collapses to the right leaf type.
6081        let run_ends = Int32Array::from(vec![2i32, 3, 5]);
6082
6083        let col_a: ArrayRef = Arc::new(
6084            RunArray::try_new(
6085                &run_ends,
6086                &StringArray::from(vec![Some("foo"), None, Some("bar")]),
6087            )
6088            .unwrap(),
6089        );
6090        let col_b: ArrayRef = Arc::new(
6091            RunArray::try_new(&run_ends, &Int32Array::from(vec![Some(1), None, Some(2)])).unwrap(),
6092        );
6093
6094        let struct_array: ArrayRef = Arc::new(StructArray::new(
6095            Fields::from(vec![
6096                Field::new("a", col_a.data_type().clone(), true),
6097                Field::new("b", col_b.data_type().clone(), true),
6098            ]),
6099            vec![col_a, col_b],
6100            None,
6101        ));
6102
6103        let schema = Arc::new(Schema::new(vec![Field::new(
6104            "row",
6105            struct_array.data_type().clone(),
6106            true,
6107        )]));
6108        let batch = RecordBatch::try_new(schema.clone(), vec![struct_array]).unwrap();
6109
6110        let mut buf = Vec::new();
6111        let mut writer = ArrowWriter::try_new(&mut buf, schema, None).unwrap();
6112        writer.write(&batch).unwrap();
6113        let metadata = writer.close().unwrap();
6114
6115        let parquet_schema = metadata.file_metadata().schema_descr();
6116        assert_eq!(parquet_schema.num_columns(), 2);
6117        assert_eq!(
6118            parquet_schema.column(0).physical_type(),
6119            crate::basic::Type::BYTE_ARRAY
6120        );
6121        assert_eq!(parquet_schema.column(0).path().string(), "row.a");
6122        assert_eq!(
6123            parquet_schema.column(1).physical_type(),
6124            crate::basic::Type::INT32
6125        );
6126        assert_eq!(parquet_schema.column(1).path().string(), "row.b");
6127    }
6128}