Skip to main content

parquet/arrow/arrow_writer/
mod.rs

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