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, compute::cast};
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    /// Writes each of `arrays` as a batch, in several mini-batches, checks they round trip,
2626    /// and returns the number of values in the column's dictionary page
2627    fn write_string_dictionaries(arrays: &[ArrayRef]) -> usize {
2628        let schema = Arc::new(Schema::new(vec![Field::new(
2629            "col",
2630            arrays[0].data_type().clone(),
2631            true,
2632        )]));
2633        let props = WriterProperties::builder().set_write_batch_size(2).build();
2634        let mut file = vec![];
2635        let mut writer = ArrowWriter::try_new(&mut file, schema.clone(), Some(props)).unwrap();
2636        for array in arrays {
2637            let batch = RecordBatch::try_new(schema.clone(), vec![array.clone()]).unwrap();
2638            writer.write(&batch).unwrap();
2639        }
2640        writer.close().unwrap();
2641
2642        let reader = SerializedFileReader::new(Bytes::from(file.clone())).unwrap();
2643        let mut pages = reader
2644            .get_row_group(0)
2645            .unwrap()
2646            .get_column_page_reader(0)
2647            .unwrap();
2648        let dictionary_page = pages.get_next_page().unwrap().unwrap();
2649        assert!(matches!(dictionary_page, Page::DictionaryPage { .. }));
2650
2651        let expected: Vec<&dyn Array> = arrays.iter().map(|a| a.as_ref()).collect();
2652        let expected = concat(&expected).unwrap();
2653        let actual = read_column(file);
2654        assert_eq!(
2655            cast(&actual, &DataType::Utf8).unwrap().as_ref(),
2656            cast(&expected, &DataType::Utf8).unwrap().as_ref()
2657        );
2658
2659        dictionary_page.num_values() as usize
2660    }
2661
2662    #[test]
2663    fn arrow_writer_string_dictionary_cache_reused_for_same_values() {
2664        // "dup" is in the dictionary twice, and values are both short enough to be inlined and not
2665        let shared: ArrayRef = Arc::new(StringArray::from(vec![
2666            "a",
2667            "bb",
2668            "longer than eight bytes",
2669            "dup",
2670            "dup",
2671        ]));
2672        let keys = Int32Array::from(vec![
2673            Some(0),
2674            Some(1),
2675            None,
2676            Some(2),
2677            Some(3),
2678            Some(4),
2679            Some(0),
2680        ]);
2681        let first = DictionaryArray::new(keys, shared.clone());
2682        // Same values `Arc` as `first`, so encoded with the cache built for `first`. Sliced,
2683        // with fewer rows than values, so the cache is only used because it is reused
2684        let keys = Int32Array::from(vec![Some(4), Some(4), Some(1), None, Some(2), Some(0)]);
2685        let second = DictionaryArray::new(keys, shared).slice(1, 4);
2686
2687        let arrays: Vec<ArrayRef> = vec![Arc::new(first), Arc::new(second)];
2688        // a, bb, longer than eight bytes, dup
2689        assert_eq!(write_string_dictionaries(&arrays), 4);
2690    }
2691
2692    #[test]
2693    fn arrow_writer_string_dictionary_cache_replaced_for_new_values() {
2694        let first_values: ArrayRef = Arc::new(StringArray::from(vec!["a", "bb", "ccc"]));
2695        let first = DictionaryArray::new(Int32Array::from(vec![0, 1, 2, 0]), first_values.clone());
2696        // A different values `Arc`, sharing some values with `first` at other keys, so a
2697        // stale cache would map keys to the wrong values
2698        let second_values: ArrayRef = Arc::new(StringArray::from(vec!["zz", "a", "new value!!"]));
2699        let second = DictionaryArray::new(
2700            Int32Array::from(vec![Some(2), Some(1), None, Some(0), Some(1)]),
2701            second_values,
2702        );
2703        // Back to `first`'s values, whose cache has since been replaced
2704        let third = DictionaryArray::new(Int32Array::from(vec![2, 1, 0]), first_values);
2705
2706        let arrays: Vec<ArrayRef> = vec![Arc::new(first), Arc::new(second), Arc::new(third)];
2707        // a, bb, ccc, zz, new value!!
2708        assert_eq!(write_string_dictionaries(&arrays), 5);
2709    }
2710
2711    #[test]
2712    fn roundtrip_bytes() {
2713        // define schema
2714        let schema = Arc::new(Schema::new(vec![
2715            Field::new("a", DataType::Int32, false),
2716            Field::new("b", DataType::Int32, true),
2717        ]));
2718
2719        // create some data
2720        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2721        let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2722
2723        // build a record batch
2724        let expected_batch =
2725            RecordBatch::try_new(schema.clone(), vec![Arc::new(a), Arc::new(b)]).unwrap();
2726
2727        for buffer in [
2728            get_bytes_after_close(schema.clone(), &expected_batch),
2729            get_bytes_by_into_inner(schema, &expected_batch),
2730        ] {
2731            let cursor = Bytes::from(buffer);
2732            let mut record_batch_reader = ParquetRecordBatchReader::try_new(cursor, 1024).unwrap();
2733
2734            let actual_batch = record_batch_reader
2735                .next()
2736                .expect("No batch found")
2737                .expect("Unable to get batch");
2738
2739            assert_eq!(expected_batch.schema(), actual_batch.schema());
2740            assert_eq!(expected_batch.num_columns(), actual_batch.num_columns());
2741            assert_eq!(expected_batch.num_rows(), actual_batch.num_rows());
2742            for i in 0..expected_batch.num_columns() {
2743                let expected_data = expected_batch.column(i).to_data();
2744                let actual_data = actual_batch.column(i).to_data();
2745
2746                assert_eq!(expected_data, actual_data);
2747            }
2748        }
2749    }
2750
2751    fn read_column(file: Vec<u8>) -> ArrayRef {
2752        let reader = ParquetRecordBatchReader::try_new(Bytes::from(file), 4096).unwrap();
2753        let batches: Vec<RecordBatch> = reader.map(|b| b.unwrap()).collect();
2754        let arrays: Vec<&dyn Array> = batches.iter().map(|b| b.column(0).as_ref()).collect();
2755        concat(&arrays).unwrap()
2756    }
2757
2758    fn roundtrip_compatible_column(field: Field, col: ArrayRef) -> ArrayRef {
2759        let writer_schema = Arc::new(Schema::new(vec![field]));
2760        let batch_schema = Arc::new(Schema::new(vec![Field::new(
2761            "c",
2762            col.data_type().clone(),
2763            col.logical_null_count() != 0,
2764        )]));
2765        let batch = RecordBatch::try_new(batch_schema, vec![col]).unwrap();
2766        // Isolate level/value compatibility from Arrow metadata reconstruction,
2767        // which is covered separately by the default-reader regression below.
2768        let options = ArrowWriterOptions::new().with_skip_arrow_metadata(true);
2769        let mut file = vec![];
2770        let mut writer =
2771            ArrowWriter::try_new_with_options(&mut file, writer_schema, options).unwrap();
2772        writer.write(&batch).unwrap();
2773        writer.close().unwrap();
2774        read_column(file)
2775    }
2776
2777    fn assert_compatible_roundtrip(field: Field, actual: ArrayRef, expected: &ArrayRef) {
2778        assert!(!compute_leaves(&field, &actual).unwrap().is_empty());
2779        assert_eq!(
2780            roundtrip_compatible_column(field, actual).as_ref(),
2781            expected.as_ref()
2782        );
2783    }
2784
2785    fn run_end_type(value_type: DataType, nullable_values: bool) -> DataType {
2786        DataType::RunEndEncoded(
2787            Arc::new(Field::new(
2788                Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
2789                DataType::Int32,
2790                false,
2791            )),
2792            Arc::new(Field::new(
2793                Field::REE_VALUES_FIELD_DEFAULT_NAME,
2794                value_type,
2795                nullable_values,
2796            )),
2797        )
2798    }
2799
2800    #[test]
2801    fn arrow_writer_required_ree_schema_rejects_nullable_dense_input() {
2802        let writer_schema = Arc::new(Schema::new(vec![Field::new(
2803            "c",
2804            run_end_type(DataType::Int32, false),
2805            false,
2806        )]));
2807        let batch_schema = Arc::new(Schema::new(vec![Field::new("c", DataType::Int32, true)]));
2808        let batch = RecordBatch::try_new(
2809            batch_schema.clone(),
2810            vec![Arc::new(Int32Array::from(vec![Some(7), None, Some(9)]))],
2811        )
2812        .unwrap();
2813
2814        for skip_metadata in [false, true] {
2815            for cdc in [false, true] {
2816                let props = WriterProperties::builder()
2817                    .set_content_defined_chunking(cdc.then(Default::default))
2818                    .build();
2819                let options = ArrowWriterOptions::new()
2820                    .with_properties(props)
2821                    .with_skip_arrow_metadata(skip_metadata);
2822                let mut writer = ArrowWriter::try_new_with_options(
2823                    Vec::new(),
2824                    writer_schema.clone(),
2825                    options.clone(),
2826                )
2827                .unwrap();
2828                let error = writer.write(&batch).unwrap_err().to_string();
2829                assert!(
2830                    error.contains("Found null at index 1 for required field 'c'"),
2831                    "skip_metadata={skip_metadata}, cdc={cdc}: {error}"
2832                );
2833
2834                // A nullable batch schema is fine when the data has no actual nulls.
2835                let values: ArrayRef = Arc::new(Int32Array::from(vec![7, 8, 9]));
2836                let batch =
2837                    RecordBatch::try_new(batch_schema.clone(), vec![values.clone()]).unwrap();
2838                let mut file = vec![];
2839                let mut writer =
2840                    ArrowWriter::try_new_with_options(&mut file, writer_schema.clone(), options)
2841                        .unwrap();
2842                writer.write(&batch).unwrap();
2843                writer.close().unwrap();
2844                assert_eq!(read_column(file).as_ref(), values.as_ref());
2845            }
2846        }
2847    }
2848
2849    #[test]
2850    fn ree_nulls_respect_dense_schema_nullability() {
2851        let run_ends = Int32Array::from(vec![2, 4, 5]);
2852        let values = Int32Array::from(vec![Some(1), None, Some(2)]);
2853        let ree: ArrayRef = Arc::new(Int32RunArray::try_new(&run_ends, &values).unwrap());
2854        let ree_batch = RecordBatch::try_new(
2855            Arc::new(Schema::new(vec![Field::new(
2856                "c",
2857                ree.data_type().clone(),
2858                true,
2859            )])),
2860            vec![ree],
2861        )
2862        .unwrap();
2863
2864        let nullable_schema = Arc::new(Schema::new(vec![Field::new("c", DataType::Int32, true)]));
2865        let mut file = vec![];
2866        let mut writer = ArrowWriter::try_new(&mut file, nullable_schema, None).unwrap();
2867        writer.write(&ree_batch).unwrap();
2868        writer.close().unwrap();
2869        assert_eq!(
2870            read_column(file).as_primitive::<Int32Type>(),
2871            &Int32Array::from(vec![Some(1), Some(1), None, None, Some(2)])
2872        );
2873
2874        let required_schema = Arc::new(Schema::new(vec![Field::new("c", DataType::Int32, false)]));
2875        let mut writer = ArrowWriter::try_new(Vec::new(), required_schema, None).unwrap();
2876        let error = writer.write(&ree_batch).unwrap_err().to_string();
2877        assert!(error.contains("Found null") && error.contains("required field 'c'"));
2878    }
2879
2880    #[test]
2881    fn required_null_validation_respects_reachability_and_write_timing() {
2882        let required = Field::new("c", DataType::Int32, false);
2883        let nullable: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None]));
2884        let leaves = compute_leaves(&required, &nullable).unwrap();
2885        let schema = Arc::new(Schema::new(vec![required.clone()]));
2886        let writer = ArrowWriter::try_new(Vec::new(), schema, None).unwrap();
2887        let (_file, factory) = writer.into_serialized_writer().unwrap();
2888        let mut columns = factory.create_column_writers(0).unwrap();
2889        assert!(
2890            columns[0]
2891                .write(&leaves[0])
2892                .unwrap_err()
2893                .to_string()
2894                .contains("required field 'c'")
2895        );
2896
2897        // Only a referenced logical dictionary null is invalid, not an unused value.
2898        let values: ArrayRef = Arc::new(Int32Array::from(vec![Some(7), None]));
2899        let dict: ArrayRef = Arc::new(DictionaryArray::new(
2900            Int8Array::from(vec![0, 0]),
2901            values.clone(),
2902        ));
2903        assert_eq!(
2904            roundtrip_compatible_column(required.clone(), dict).as_primitive::<Int32Type>(),
2905            &Int32Array::from(vec![7, 7])
2906        );
2907        let dict: ArrayRef = Arc::new(DictionaryArray::new(Int8Array::from(vec![0, 1]), values));
2908        let leaves = compute_leaves(&required, &dict).unwrap();
2909        assert!(leaves[0].0.validate().is_err());
2910        let dict: ArrayRef = Arc::new(DictionaryArray::new(
2911            Int8Array::from(vec![Some(0), None]),
2912            Arc::new(Int32Array::from(vec![7])),
2913        ));
2914        let leaves = compute_leaves(&required, &dict).unwrap();
2915        assert!(leaves[0].0.validate().is_err());
2916
2917        // A nullable ancestor masks an otherwise required child null, including
2918        // when the target child has a schema-only REE wrapper.
2919        let actual_fields = Fields::from(vec![Field::new("source_child", DataType::Int32, true)]);
2920        let array: ArrayRef = Arc::new(StructArray::new(
2921            actual_fields.clone(),
2922            vec![nullable.clone()],
2923            Some(NullBuffer::from(vec![true, false])),
2924        ));
2925        let unmasked: ArrayRef = Arc::new(StructArray::new(
2926            actual_fields,
2927            vec![nullable.clone()],
2928            None,
2929        ));
2930        for child_type in [DataType::Int32, run_end_type(DataType::Int32, false)] {
2931            let target = Field::new(
2932                "c",
2933                DataType::Struct(Fields::from(vec![Field::new(
2934                    "target_child",
2935                    child_type,
2936                    false,
2937                )])),
2938                true,
2939            );
2940            let leaves = compute_leaves(&target, &unmasked).unwrap();
2941            assert!(leaves[0].0.validate().is_err());
2942            let leaves = compute_leaves(&target, &array).unwrap();
2943            assert!(leaves[0].0.validate().is_ok());
2944            let result = roundtrip_compatible_column(target, array.clone());
2945            assert_eq!(result.as_struct().fields()[0].name(), "target_child");
2946            assert!(result.is_null(1));
2947        }
2948
2949        // Unused list child slots and out-of-slice scalar nulls are unreachable.
2950        let values: ArrayRef = Arc::new(Int32Array::from(vec![None, Some(9), None]));
2951        let list: ArrayRef = Arc::new(ListArray::new(
2952            Arc::new(Field::new_list_field(DataType::Int32, true)),
2953            OffsetBuffer::new(vec![1_i32, 2].into()),
2954            values,
2955            None,
2956        ));
2957        let target = Field::new(
2958            "c",
2959            DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false))),
2960            false,
2961        );
2962        assert!(
2963            compute_leaves(&target, &list).unwrap()[0]
2964                .0
2965                .validate()
2966                .is_ok()
2967        );
2968        assert_eq!(
2969            roundtrip_compatible_column(required, nullable.slice(0, 1)).as_primitive::<Int32Type>(),
2970            &Int32Array::from(vec![1])
2971        );
2972    }
2973
2974    fn check_required_list_validation(list_type: fn(Arc<Field>) -> DataType) {
2975        let make_array = |values: Vec<Option<i32>>,
2976                          offsets: Vec<i32>,
2977                          sizes: Vec<i32>,
2978                          nulls: Option<NullBuffer>| {
2979            let child = Arc::new(Field::new("source_child", DataType::Int32, true));
2980            let values: ArrayRef = Arc::new(Int32Array::from(values));
2981            match list_type(child.clone()) {
2982                DataType::List(_) => Arc::new(ListArray::new(
2983                    child,
2984                    OffsetBuffer::new(offsets.into()),
2985                    values,
2986                    nulls,
2987                )) as ArrayRef,
2988                DataType::LargeList(_) => Arc::new(LargeListArray::new(
2989                    child,
2990                    OffsetBuffer::new(
2991                        offsets
2992                            .into_iter()
2993                            .map(i64::from)
2994                            .collect::<Vec<_>>()
2995                            .into(),
2996                    ),
2997                    values,
2998                    nulls,
2999                )) as ArrayRef,
3000                DataType::ListView(_) => Arc::new(ListViewArray::new(
3001                    child,
3002                    offsets.into(),
3003                    sizes.into(),
3004                    values,
3005                    nulls,
3006                )) as ArrayRef,
3007                DataType::LargeListView(_) => Arc::new(LargeListViewArray::new(
3008                    child,
3009                    offsets
3010                        .into_iter()
3011                        .map(i64::from)
3012                        .collect::<Vec<_>>()
3013                        .into(),
3014                    sizes.into_iter().map(i64::from).collect::<Vec<_>>().into(),
3015                    values,
3016                    nulls,
3017                )) as ArrayRef,
3018                _ => unreachable!(),
3019            }
3020        };
3021        let is_view = matches!(
3022            list_type(Arc::new(Field::new("item", DataType::Int32, true))),
3023            DataType::ListView(_) | DataType::LargeListView(_)
3024        );
3025        let offsets = |starts: Vec<i32>, end: i32| {
3026            let mut result = starts;
3027            if !is_view {
3028                result.push(end);
3029            }
3030            result
3031        };
3032        let target = Field::new(
3033            "c",
3034            list_type(Arc::new(Field::new("target_child", DataType::Int32, false))),
3035            true,
3036        );
3037        let invalid = make_array(vec![None, Some(7)], offsets(vec![0], 2), vec![2], None);
3038        let valid = make_array(vec![Some(7), Some(8)], offsets(vec![0], 2), vec![2], None);
3039        let masked = make_array(
3040            vec![None, Some(7)],
3041            offsets(vec![0, 1], 2),
3042            vec![1, 1],
3043            Some(NullBuffer::from(vec![false, true])),
3044        );
3045        let unused = make_array(
3046            vec![None, Some(7), None],
3047            offsets(vec![1], 2),
3048            vec![1],
3049            None,
3050        );
3051        let empty = make_array(vec![None], offsets(vec![0], 0), vec![0], None);
3052        let sliced = make_array(
3053            vec![None, Some(7)],
3054            offsets(vec![0, 1], 2),
3055            vec![1, 1],
3056            None,
3057        )
3058        .slice(1, 1);
3059        let schema = Arc::new(Schema::new(vec![target.clone()]));
3060
3061        // Binding succeeds even for invalid input; writing the affected leaf rejects it.
3062        for (name, array, invalid) in [
3063            ("reachable null", invalid, true),
3064            ("valid required child", valid.clone(), false),
3065            ("null ancestor", masked, false),
3066            ("unused child slots", unused, false),
3067            ("empty list", empty, false),
3068            ("out-of-slice null", sliced, false),
3069        ] {
3070            let leaves = compute_leaves(&target, &array).unwrap();
3071            assert_eq!(leaves.len(), 1);
3072            for cdc in [false, true] {
3073                let props = WriterProperties::builder()
3074                    .set_content_defined_chunking(cdc.then(Default::default))
3075                    .build();
3076                let mut writer =
3077                    ArrowWriter::try_new(Vec::new(), schema.clone(), Some(props)).unwrap();
3078                let mut chunkers = writer.cdc_chunkers.take();
3079                let (_file, factory) = writer.into_serialized_writer().unwrap();
3080                let mut columns = factory.create_column_writers(0).unwrap();
3081                let before = (
3082                    columns[0].memory_size(),
3083                    columns[0].get_estimated_total_bytes(),
3084                );
3085                let result = match &mut chunkers {
3086                    Some(chunkers) => columns[0].write_with_chunker(&leaves[0], &mut chunkers[0]),
3087                    None => columns[0].write(&leaves[0]),
3088                };
3089                if invalid {
3090                    let error = result.unwrap_err().to_string();
3091                    assert!(
3092                        error.contains("Found null at index 0 for required field 'target_child'"),
3093                        "{name}, cdc={cdc}: {error}"
3094                    );
3095                    assert_eq!(
3096                        before,
3097                        (
3098                            columns[0].memory_size(),
3099                            columns[0].get_estimated_total_bytes()
3100                        ),
3101                        "invalid leaf must be rejected before encoding"
3102                    );
3103                    // A rejected leaf write must leave no encoded values behind.
3104                    let valid_leaves = compute_leaves(&target, &valid).unwrap();
3105                    match &mut chunkers {
3106                        Some(chunkers) => columns[0]
3107                            .write_with_chunker(&valid_leaves[0], &mut chunkers[0])
3108                            .unwrap(),
3109                        None => columns[0].write(&valid_leaves[0]).unwrap(),
3110                    }
3111                    assert_eq!(
3112                        columns
3113                            .pop()
3114                            .unwrap()
3115                            .close()
3116                            .unwrap()
3117                            .close
3118                            .metadata
3119                            .num_values(),
3120                        2
3121                    );
3122                } else {
3123                    result.unwrap();
3124                    columns.pop().unwrap().close().unwrap();
3125                }
3126
3127                // Exercise the ordinary and CDC row-group driving entry points as well.
3128                let batch = RecordBatch::try_new(
3129                    Arc::new(Schema::new(vec![Field::new(
3130                        "c",
3131                        array.data_type().clone(),
3132                        true,
3133                    )])),
3134                    vec![array.clone()],
3135                )
3136                .unwrap();
3137                let props = WriterProperties::builder()
3138                    .set_content_defined_chunking(cdc.then(Default::default))
3139                    .build();
3140                let mut writer =
3141                    ArrowWriter::try_new(Vec::new(), schema.clone(), Some(props)).unwrap();
3142                let result = writer.write(&batch);
3143                if invalid {
3144                    assert!(
3145                        result
3146                            .unwrap_err()
3147                            .to_string()
3148                            .contains("required field 'target_child'")
3149                    );
3150                } else {
3151                    result.unwrap();
3152                    writer.close().unwrap();
3153                }
3154            }
3155        }
3156    }
3157
3158    #[test]
3159    fn required_list_validation_preserves_level_generation() {
3160        check_required_list_validation(DataType::List);
3161    }
3162
3163    #[test]
3164    fn required_large_list_validation_preserves_level_generation() {
3165        check_required_list_validation(DataType::LargeList);
3166    }
3167
3168    #[test]
3169    fn required_list_view_validation_preserves_level_generation() {
3170        check_required_list_validation(DataType::ListView);
3171    }
3172
3173    #[test]
3174    fn required_large_list_view_validation_preserves_level_generation() {
3175        check_required_list_validation(DataType::LargeListView);
3176    }
3177
3178    fn nullable_int32_list() -> ArrayRef {
3179        Arc::new(ListArray::new(
3180            Arc::new(Field::new_list_field(DataType::Int32, true)),
3181            OffsetBuffer::new(vec![0_i32, 2, 3].into()),
3182            Arc::new(Int32Array::from(vec![Some(1), None, Some(3)])),
3183            None,
3184        ))
3185    }
3186
3187    #[test]
3188    fn arrow_writer_dense_list_under_ree_schema() {
3189        let list = nullable_int32_list();
3190        // The schema wrapper is physical, not a second logical list node.
3191        let field = Field::new("c", run_end_type(list.data_type().clone(), false), false);
3192        assert_compatible_roundtrip(field, list.clone(), &list);
3193    }
3194
3195    #[test]
3196    fn arrow_writer_dense_list_under_dictionary_schema() {
3197        let list = nullable_int32_list();
3198        let field = Field::new(
3199            "c",
3200            DataType::Dictionary(Box::new(DataType::Int8), Box::new(list.data_type().clone())),
3201            false,
3202        );
3203        assert_compatible_roundtrip(field, list.clone(), &list);
3204    }
3205
3206    #[test]
3207    fn arrow_writer_dense_values_under_nested_ree_schema() {
3208        let dense: ArrayRef = Arc::new(Int32Array::from(vec![Some(7), None, Some(9)]));
3209        let field = Field::new(
3210            "c",
3211            run_end_type(run_end_type(DataType::Int32, true), false),
3212            false,
3213        );
3214        assert_compatible_roundtrip(field, dense.clone(), &dense);
3215    }
3216
3217    #[test]
3218    fn arrow_writer_dense_values_under_dictionary_ree_schema() {
3219        let dense: ArrayRef = Arc::new(Int32Array::from(vec![Some(7), None, Some(9)]));
3220        let field = Field::new(
3221            "c",
3222            DataType::Dictionary(
3223                Box::new(DataType::Int8),
3224                Box::new(run_end_type(DataType::Int32, true)),
3225            ),
3226            false,
3227        );
3228        assert_compatible_roundtrip(field, dense.clone(), &dense);
3229    }
3230
3231    #[test]
3232    fn arrow_writer_dense_values_under_ree_dictionary_schema() {
3233        let dense: ArrayRef = Arc::new(Int32Array::from(vec![Some(7), None, Some(9)]));
3234        // Hoist REE value nullability before peeling the dictionary; the
3235        // dense batch contains an actual null despite the required outer field.
3236        let field = Field::new(
3237            "c",
3238            run_end_type(
3239                DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int32)),
3240                true,
3241            ),
3242            false,
3243        );
3244        assert_compatible_roundtrip(field, dense.clone(), &dense);
3245    }
3246
3247    #[test]
3248    fn arrow_writer_dense_struct_under_ree_child_schema() {
3249        let dense: ArrayRef = Arc::new(Int32Array::from(vec![Some(7), None, Some(9)]));
3250        let struct_fields = Fields::from(vec![Field::new("value", DataType::Int32, true)]);
3251        let dense_struct: ArrayRef = Arc::new(StructArray::new(struct_fields, vec![dense], None));
3252        // A wrapper on a child must not make the enclosing struct incompatible.
3253        let field = Field::new(
3254            "c",
3255            DataType::Struct(Fields::from(vec![Field::new(
3256                "value",
3257                run_end_type(DataType::Int32, true),
3258                false,
3259            )])),
3260            false,
3261        );
3262        assert_compatible_roundtrip(field, dense_struct.clone(), &dense_struct);
3263    }
3264
3265    #[test]
3266    fn arrow_writer_dense_list_under_ree_child_schema() {
3267        let list = nullable_int32_list();
3268        let field = Field::new(
3269            "c",
3270            DataType::List(Arc::new(Field::new(
3271                "item",
3272                run_end_type(DataType::Int32, true),
3273                false,
3274            ))),
3275            false,
3276        );
3277        assert_compatible_roundtrip(field, list.clone(), &list);
3278    }
3279
3280    #[test]
3281    fn arrow_writer_dense_list_under_dictionary_child_schema() {
3282        let list = nullable_int32_list();
3283        let field = Field::new(
3284            "c",
3285            DataType::List(Arc::new(Field::new_list_field(
3286                DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int32)),
3287                true,
3288            ))),
3289            false,
3290        );
3291        assert_compatible_roundtrip(field, list.clone(), &list);
3292    }
3293
3294    #[test]
3295    fn arrow_writer_dictionary_list_under_dense_child_schema() {
3296        let dict: ArrayRef = Arc::new(DictionaryArray::new(
3297            Int8Array::from(vec![Some(0), None, Some(1)]),
3298            Arc::new(Int32Array::from(vec![1, 3])),
3299        ));
3300        let dict_list: ArrayRef = Arc::new(ListArray::new(
3301            Arc::new(Field::new_list_field(dict.data_type().clone(), true)),
3302            OffsetBuffer::new(vec![0_i32, 2, 3].into()),
3303            dict,
3304            None,
3305        ));
3306        let expected = nullable_int32_list();
3307        assert_compatible_roundtrip(
3308            Field::new("c", expected.data_type().clone(), false),
3309            dict_list,
3310            &expected,
3311        );
3312    }
3313
3314    #[test]
3315    fn arrow_writer_dense_map_under_ree_value_schema() {
3316        let key_field = Arc::new(Field::new("keys", DataType::Utf8, false));
3317        let entries = StructArray::new(
3318            Fields::from(vec![
3319                key_field.clone(),
3320                Arc::new(Field::new("values", DataType::Int32, true)),
3321            ]),
3322            vec![
3323                Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef,
3324                Arc::new(Int32Array::from(vec![Some(7), None, Some(9)])),
3325            ],
3326            None,
3327        );
3328        let map: ArrayRef = Arc::new(MapArray::new(
3329            Arc::new(Field::new("entries", entries.data_type().clone(), false)),
3330            OffsetBuffer::new(vec![0_i32, 2, 3].into()),
3331            entries,
3332            None,
3333            false,
3334        ));
3335        // Resolve the value wrapper after walking through the entries struct.
3336        let field = Field::new(
3337            "c",
3338            DataType::Map(
3339                Arc::new(Field::new(
3340                    "entries",
3341                    DataType::Struct(Fields::from(vec![
3342                        key_field,
3343                        Arc::new(Field::new(
3344                            "values",
3345                            run_end_type(DataType::Int32, true),
3346                            false,
3347                        )),
3348                    ])),
3349                    false,
3350                )),
3351                false,
3352            ),
3353            false,
3354        );
3355        assert_compatible_roundtrip(field, map.clone(), &map);
3356    }
3357
3358    #[test]
3359    fn arrow_writer_nested_dictionary_metadata_roundtrips_with_default_reader() {
3360        let dense: ArrayRef = Arc::new(StructArray::new(
3361            Fields::from(vec![Field::new("item", DataType::Int32, true)]),
3362            vec![Arc::new(Int32Array::from(vec![Some(7), None, Some(9)]))],
3363            None,
3364        ));
3365        let field = Field::new(
3366            "value",
3367            DataType::Dictionary(
3368                Box::new(DataType::Int8),
3369                Box::new(dense.data_type().clone()),
3370            ),
3371            false,
3372        )
3373        .with_metadata(HashMap::from([(
3374            "description".to_string(),
3375            "nested dictionary".to_string(),
3376        )]));
3377        let writer_schema = Arc::new(Schema::new(vec![field.clone()]));
3378        let expected_schema = Arc::new(Schema::new(vec![
3379            field.with_data_type(dense.data_type().clone()),
3380        ]));
3381        let batch = RecordBatch::try_new(expected_schema.clone(), vec![dense]).unwrap();
3382
3383        // Keep Arrow metadata enabled: ignoring it would hide the original
3384        // reader-initialization failure caused by the nested dictionary hint.
3385        let mut file = vec![];
3386        let mut writer = ArrowWriter::try_new(&mut file, writer_schema, None).unwrap();
3387        writer.write(&batch).unwrap();
3388        writer.close().unwrap();
3389
3390        let builder = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(file)).unwrap();
3391        assert!(
3392            builder
3393                .metadata()
3394                .file_metadata()
3395                .key_value_metadata()
3396                .unwrap()
3397                .iter()
3398                .any(|kv| kv.key == ARROW_SCHEMA_META_KEY)
3399        );
3400        assert_eq!(builder.schema(), &expected_schema);
3401        let decoded = builder
3402            .build()
3403            .unwrap()
3404            .collect::<ArrowResult<Vec<_>>>()
3405            .unwrap();
3406        assert_eq!(decoded, vec![batch]);
3407    }
3408
3409    #[test]
3410    fn arrow_writer_rejects_missing_or_extra_top_level_columns_cleanly() {
3411        let one_field = Arc::new(Schema::new(vec![Field::new(
3412            "a",
3413            ArrowDataType::Int32,
3414            false,
3415        )]));
3416        let two_fields = Arc::new(Schema::new(vec![
3417            Field::new("a", ArrowDataType::Int32, false),
3418            Field::new("b", ArrowDataType::Int32, false),
3419        ]));
3420        let one_column = RecordBatch::try_new(
3421            Arc::clone(&one_field),
3422            vec![Arc::new(Int32Array::from(vec![1]))],
3423        )
3424        .unwrap();
3425        let two_columns = RecordBatch::try_new(
3426            Arc::clone(&two_fields),
3427            vec![
3428                Arc::new(Int32Array::from(vec![1])),
3429                Arc::new(Int32Array::from(vec![2])),
3430            ],
3431        )
3432        .unwrap();
3433
3434        let mut writer = ArrowWriter::try_new(Vec::new(), Arc::clone(&two_fields), None).unwrap();
3435        let err = writer.write(&one_column).unwrap_err();
3436        assert!(
3437            err.to_string()
3438                .contains("writer has 2 top-level fields but batch has 1 columns"),
3439            "{err}"
3440        );
3441        assert_eq!(writer.in_progress_rows(), 0);
3442
3443        let mut writer = ArrowWriter::try_new(Vec::new(), Arc::clone(&one_field), None).unwrap();
3444        let err = writer.write(&two_columns).unwrap_err();
3445        assert!(
3446            err.to_string()
3447                .contains("writer has 1 top-level fields but batch has 2 columns"),
3448            "{err}"
3449        );
3450        assert_eq!(writer.in_progress_rows(), 0);
3451    }
3452
3453    #[test]
3454    fn arrow_row_group_factory_rejects_parquet_leaf_count_mismatch_cleanly() {
3455        let one_field = Arc::new(Schema::new(vec![Field::new(
3456            "a",
3457            ArrowDataType::Int32,
3458            false,
3459        )]));
3460        let two_fields = Arc::new(Schema::new(vec![
3461            Field::new("a", ArrowDataType::Int32, false),
3462            Field::new("b", ArrowDataType::Int32, false),
3463        ]));
3464        let parquet = ArrowSchemaConverter::new().convert(&one_field).unwrap();
3465        let props = Arc::new(WriterProperties::default());
3466        let file_writer =
3467            SerializedFileWriter::new(Vec::new(), parquet.root_schema_ptr(), Arc::clone(&props))
3468                .unwrap();
3469
3470        let eager_error =
3471            ArrowRowGroupWriterFactory::try_new(&file_writer, Arc::clone(&two_fields)).unwrap_err();
3472        assert!(matches!(&eager_error, ParquetError::General(_)));
3473        #[expect(deprecated)] // Verify that the legacy constructor still defers errors.
3474        let factory = ArrowRowGroupWriterFactory::new(&file_writer, Arc::clone(&two_fields));
3475        let Err(err) = factory.create_column_writers(0) else {
3476            panic!("mismatched schemas should not create column writers")
3477        };
3478        assert_eq!(err.to_string(), eager_error.to_string());
3479        assert!(
3480            err.to_string()
3481                .contains("Arrow schema maps to 2 leaf columns but Parquet schema has 1"),
3482            "{err}"
3483        );
3484
3485        let parquet = ArrowSchemaConverter::new().convert(&two_fields).unwrap();
3486        let file_writer =
3487            SerializedFileWriter::new(Vec::new(), parquet.root_schema_ptr(), props).unwrap();
3488        let eager_error =
3489            ArrowRowGroupWriterFactory::try_new(&file_writer, Arc::clone(&one_field)).unwrap_err();
3490        assert!(matches!(&eager_error, ParquetError::General(_)));
3491        #[expect(deprecated)] // Verify that the legacy constructor still defers errors.
3492        let factory = ArrowRowGroupWriterFactory::new(&file_writer, one_field);
3493        let Err(err) = factory.create_column_writers(0) else {
3494            panic!("mismatched schemas should not create column writers")
3495        };
3496        assert_eq!(err.to_string(), eager_error.to_string());
3497        assert!(
3498            err.to_string()
3499                .contains("Arrow schema maps to 1 leaf columns but Parquet schema has 2"),
3500            "{err}"
3501        );
3502    }
3503
3504    #[test]
3505    fn arrow_writer_required_map_keys_reject_only_reachable_nulls() {
3506        let keys: ArrayRef = Arc::new(DictionaryArray::<Int32Type>::new(
3507            Int32Array::from(vec![1, 0, 1]),
3508            Arc::new(Int32Array::from(vec![None, Some(7)])),
3509        ));
3510        let values = Int32Array::from(vec![Some(10), Some(20), None]);
3511        let fields = Fields::from(vec![
3512            Field::new("key", keys.data_type().clone(), false),
3513            Field::new("value", DataType::Int32, true),
3514        ]);
3515        assert!(
3516            StructArray::try_new(
3517                fields.clone(),
3518                vec![keys.clone(), Arc::new(values.clone())],
3519                None
3520            )
3521            .is_err()
3522        );
3523
3524        // ArrayData validates physical nulls, not null dictionary values. This
3525        // safe construction reaches the writer with a logically null map key.
3526        let data = ArrayData::builder(DataType::Struct(fields))
3527            .len(3)
3528            .add_child_data(keys.to_data())
3529            .add_child_data(values.to_data())
3530            .build()
3531            .unwrap();
3532        data.validate_full().unwrap();
3533        let entries = StructArray::from(data);
3534        let entries_field = Arc::new(Field::new("entries", entries.data_type().clone(), false));
3535        let offsets = OffsetBuffer::new(vec![0i32, 1, 2, 3].into());
3536        let map = MapArray::try_new(
3537            entries_field.clone(),
3538            offsets.clone(),
3539            entries.clone(),
3540            None,
3541            false,
3542        )
3543        .unwrap();
3544        assert_eq!(map.keys().logical_null_count(), 1);
3545        let masked = MapArray::try_new(
3546            entries_field,
3547            offsets,
3548            entries,
3549            Some(NullBuffer::from(vec![true, false, true])),
3550            false,
3551        )
3552        .unwrap();
3553        let writer_schema = Arc::new(Schema::new(vec![Field::new(
3554            "map",
3555            DataType::Map(
3556                Arc::new(Field::new(
3557                    "entries",
3558                    DataType::Struct(
3559                        vec![
3560                            Field::new("key", DataType::Int32, false),
3561                            Field::new("value", DataType::Int32, true),
3562                        ]
3563                        .into(),
3564                    ),
3565                    false,
3566                )),
3567                false,
3568            ),
3569            true,
3570        )]));
3571        for cdc in [false, true] {
3572            let props = WriterProperties::builder()
3573                .set_content_defined_chunking(cdc.then(Default::default))
3574                .build();
3575            for (array, expected_keys, expected_null_maps) in [
3576                (map.clone(), None, 0),
3577                (map.slice(0, 1), Some(1), 0),
3578                (map.slice(2, 1), Some(1), 0),
3579                (masked.clone(), Some(2), 1),
3580            ] {
3581                let schema = Arc::new(Schema::new(vec![Field::new(
3582                    "map",
3583                    array.data_type().clone(),
3584                    true,
3585                )]));
3586                let batch = RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap();
3587                let mut bytes = Vec::new();
3588                let mut writer =
3589                    ArrowWriter::try_new(&mut bytes, writer_schema.clone(), Some(props.clone()))
3590                        .unwrap();
3591                let result = writer.write(&batch);
3592                let Some(expected_keys) = expected_keys else {
3593                    let err = result.unwrap_err();
3594                    assert!(
3595                        err.to_string().contains("Found null")
3596                            && err.to_string().contains("required field 'key'"),
3597                        "{err}"
3598                    );
3599                    continue;
3600                };
3601                result.unwrap();
3602                writer.close().unwrap();
3603                let actual = ParquetRecordBatchReader::try_new(Bytes::from(bytes), 1024)
3604                    .unwrap()
3605                    .next()
3606                    .unwrap()
3607                    .unwrap();
3608                let actual = actual.column(0).as_map();
3609                assert_eq!(actual.null_count(), expected_null_maps);
3610                assert_eq!(
3611                    actual.keys().as_primitive::<Int32Type>().values().as_ref(),
3612                    vec![7; expected_keys]
3613                );
3614            }
3615        }
3616    }
3617
3618    #[test]
3619    fn arrow_writer_rejects_nullable_map_key_schemas_early() {
3620        let schema_with_key = |key| {
3621            Arc::new(Schema::new(vec![Field::new(
3622                "map",
3623                DataType::Map(
3624                    Arc::new(Field::new(
3625                        "entries",
3626                        DataType::Struct(
3627                            vec![key, Field::new("value", DataType::Int32, true)].into(),
3628                        ),
3629                        false,
3630                    )),
3631                    false,
3632                ),
3633                true,
3634            )]))
3635        };
3636        let required = schema_with_key(Field::new("key", DataType::Int32, false));
3637        let parquet = ArrowSchemaConverter::new().convert(&required).unwrap();
3638        let ree = DataType::RunEndEncoded(
3639            Arc::new(Field::new("run_ends", DataType::Int32, false)),
3640            Arc::new(Field::new("values", DataType::Int32, true)),
3641        );
3642        for key in [
3643            Field::new("key", DataType::Int32, true),
3644            Field::new("key", ree.clone(), false),
3645            Field::new(
3646                "key",
3647                DataType::Dictionary(Box::new(DataType::Int8), Box::new(ree)),
3648                false,
3649            ),
3650        ] {
3651            let schema = schema_with_key(key);
3652            for explicit_schema in [false, true] {
3653                for skip_metadata in [false, true] {
3654                    let mut options =
3655                        ArrowWriterOptions::new().with_skip_arrow_metadata(skip_metadata);
3656                    if explicit_schema {
3657                        options = options.with_parquet_schema(parquet.clone());
3658                    }
3659                    let mut bytes = Vec::new();
3660                    let err =
3661                        ArrowWriter::try_new_with_options(&mut bytes, schema.clone(), options)
3662                            .unwrap_err();
3663                    assert!(
3664                        err.to_string()
3665                            .contains("Map key field 'key' must be non-nullable"),
3666                        "{err}"
3667                    );
3668                    assert!(
3669                        bytes.is_empty(),
3670                        "schema rejection must precede file initialization"
3671                    );
3672                }
3673            }
3674            let file = SerializedFileWriter::new(
3675                Vec::new(),
3676                parquet.root_schema_ptr(),
3677                Arc::new(WriterProperties::default()),
3678            )
3679            .unwrap();
3680            let eager_error =
3681                ArrowRowGroupWriterFactory::try_new(&file, Arc::clone(&schema)).unwrap_err();
3682            assert!(matches!(&eager_error, ParquetError::ArrowError(_)));
3683            assert_eq!(
3684                eager_error.to_string(),
3685                "Arrow: Map key field 'key' must be non-nullable, including its encoded values"
3686            );
3687            #[expect(deprecated)] // Verify that the legacy constructor still defers errors.
3688            let factory = ArrowRowGroupWriterFactory::new(&file, schema);
3689            for row_group_index in 0..2 {
3690                let Err(err) = factory.create_column_writers(row_group_index) else {
3691                    panic!("nullable map-key schema should not create column writers")
3692                };
3693                assert!(matches!(&err, ParquetError::ArrowError(_)), "{err:?}");
3694                assert_eq!(
3695                    err.to_string(),
3696                    "Arrow: Map key field 'key' must be non-nullable, including its encoded values"
3697                );
3698            }
3699        }
3700    }
3701
3702    #[test]
3703    fn eager_row_group_validation_rejects_writer_and_chunker_count_mismatches() {
3704        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
3705        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![1]))])
3706            .unwrap();
3707        let writer = ArrowWriter::try_new(Vec::new(), schema, None).unwrap();
3708        let (_file, factory) = writer.into_serialized_writer().unwrap();
3709        let mut group = factory.create_row_group_writer(0).unwrap();
3710        assert!(
3711            group
3712                .write_with_chunkers(&batch, &mut [])
3713                .unwrap_err()
3714                .to_string()
3715                .contains("chunkers")
3716        );
3717        assert_eq!(group.buffered_rows, 0);
3718        group.writers.clear();
3719        assert!(
3720            group
3721                .write(&batch)
3722                .unwrap_err()
3723                .to_string()
3724                .contains("column writers")
3725        );
3726        assert_eq!(group.buffered_rows, 0);
3727    }
3728
3729    #[test]
3730    fn arrow_write_schema_plan_preserves_error_variants() {
3731        for error in [
3732            ParquetError::General("schema mismatch".to_owned()),
3733            ParquetError::NYI("unsupported type".to_owned()),
3734            ParquetError::ArrowError("invalid schema".to_owned()),
3735        ] {
3736            let expected_variant = std::mem::discriminant(&error);
3737            let expected_message = error.to_string();
3738            let cached = ArrowWriteSchemaPlanError::from_parquet_error(error);
3739            for _ in 0..2 {
3740                let error = cached.to_parquet_error();
3741                assert_eq!(std::mem::discriminant(&error), expected_variant);
3742                assert_eq!(error.to_string(), expected_message);
3743            }
3744        }
3745    }
3746
3747    #[test]
3748    fn arrow_write_schema_plan_caches_nested_leaf_ranges() {
3749        let nested = ArrowDataType::Struct(Fields::from(vec![
3750            Field::new("i", ArrowDataType::Int32, false),
3751            Field::new("s", ArrowDataType::Utf8, true),
3752        ]));
3753        let list = ArrowDataType::List(Arc::new(Field::new("item", ArrowDataType::Boolean, true)));
3754        let ree = ArrowDataType::RunEndEncoded(
3755            Arc::new(Field::new("run_ends", ArrowDataType::Int32, false)),
3756            Arc::new(Field::new("values", ArrowDataType::Int64, false)),
3757        );
3758        let dictionary_fsb = ArrowDataType::Dictionary(
3759            Box::new(ArrowDataType::Int8),
3760            Box::new(ArrowDataType::FixedSizeBinary(2)),
3761        );
3762        let schema = Arc::new(Schema::new(vec![
3763            Field::new("nested", nested, true),
3764            Field::new("list", list, true),
3765            Field::new("ree", ree, false),
3766            Field::new("dictionary_fsb", dictionary_fsb, true),
3767        ]));
3768        let parquet = ArrowSchemaConverter::new().convert(&schema).unwrap();
3769
3770        let plan = ArrowWriteSchemaPlan::try_new(&parquet, &schema).unwrap();
3771        let ranges = plan
3772            .fields
3773            .iter()
3774            .map(|field| field.leaf_range.clone())
3775            .collect::<Vec<_>>();
3776        assert_eq!(ranges, vec![0..2, 2..3, 3..4, 4..5]);
3777        assert_eq!(plan.leaves.len(), parquet.num_columns());
3778        assert_eq!(
3779            plan.leaves[4].physical_type(),
3780            crate::basic::Type::FIXED_LEN_BYTE_ARRAY
3781        );
3782    }
3783
3784    #[test]
3785    fn cached_schema_plan_preserves_compatible_physical_layout_alternation() {
3786        let writer_schema = Arc::new(Schema::new(vec![
3787            Field::new("number", ArrowDataType::Int32, false),
3788            Field::new("text", ArrowDataType::Utf8, false),
3789            Field::new("bytes", ArrowDataType::Binary, false),
3790        ]));
3791
3792        let dense = RecordBatch::try_new(
3793            Arc::clone(&writer_schema),
3794            vec![
3795                Arc::new(Int32Array::from(vec![1])),
3796                Arc::new(StringArray::from(vec!["a"])),
3797                Arc::new(BinaryArray::from_iter_values([b"a".as_slice()])),
3798            ],
3799        )
3800        .unwrap();
3801
3802        let number_dictionary = DictionaryArray::<Int8Type>::try_new(
3803            Int8Array::from(vec![0]),
3804            Arc::new(Int32Array::from(vec![2])),
3805        )
3806        .unwrap();
3807        let text_dictionary = DictionaryArray::<Int8Type>::try_new(
3808            Int8Array::from(vec![0]),
3809            Arc::new(StringArray::from(vec!["b"])),
3810        )
3811        .unwrap();
3812        let bytes_dictionary = DictionaryArray::<Int8Type>::try_new(
3813            Int8Array::from(vec![0]),
3814            Arc::new(BinaryArray::from_iter_values([b"b".as_slice()])),
3815        )
3816        .unwrap();
3817        let dictionary = RecordBatch::try_from_iter(vec![
3818            ("number", Arc::new(number_dictionary) as ArrayRef),
3819            ("text", Arc::new(text_dictionary) as ArrayRef),
3820            ("bytes", Arc::new(bytes_dictionary) as ArrayRef),
3821        ])
3822        .unwrap();
3823
3824        let run_ends = Int32Array::from(vec![1]);
3825        let number_ree: ArrayRef =
3826            Arc::new(Int32RunArray::try_new(&run_ends, &Int32Array::from(vec![3])).unwrap());
3827        let text_ree: ArrayRef =
3828            Arc::new(Int32RunArray::try_new(&run_ends, &StringArray::from(vec!["c"])).unwrap());
3829        let bytes_ree: ArrayRef = Arc::new(
3830            Int32RunArray::try_new(&run_ends, &BinaryArray::from_iter_values([b"c".as_slice()]))
3831                .unwrap(),
3832        );
3833        let ree = RecordBatch::try_from_iter(vec![
3834            ("number", number_ree),
3835            ("text", text_ree),
3836            ("bytes", bytes_ree),
3837        ])
3838        .unwrap();
3839
3840        let alternate = RecordBatch::try_from_iter(vec![
3841            ("number", Arc::new(Int32Array::from(vec![4])) as ArrayRef),
3842            (
3843                "text",
3844                Arc::new(LargeStringArray::from(vec!["d"])) as ArrayRef,
3845            ),
3846            (
3847                "bytes",
3848                Arc::new(BinaryViewArray::from_iter_values([b"d".as_slice()])) as ArrayRef,
3849            ),
3850        ])
3851        .unwrap();
3852
3853        let props = WriterProperties::builder()
3854            .set_max_row_group_row_count(Some(1))
3855            .build();
3856        let mut out = Vec::new();
3857        let mut writer =
3858            ArrowWriter::try_new(&mut out, Arc::clone(&writer_schema), Some(props)).unwrap();
3859        for batch in [&dense, &dictionary, &ree, &alternate] {
3860            writer.write(batch).unwrap();
3861        }
3862        writer.close().unwrap();
3863
3864        let builder = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(out)).unwrap();
3865        assert_eq!(builder.metadata().num_row_groups(), 4);
3866        let mut reader = builder.build().unwrap();
3867        let actual = reader.next().unwrap().unwrap();
3868        assert_eq!(
3869            actual.column(0).as_primitive::<Int32Type>(),
3870            &Int32Array::from(vec![1, 2, 3, 4])
3871        );
3872        assert_eq!(
3873            actual.column(1).as_string::<i32>(),
3874            &StringArray::from(vec!["a", "b", "c", "d"])
3875        );
3876        assert_eq!(
3877            actual.column(2).as_binary::<i32>(),
3878            &BinaryArray::from_iter_values([
3879                b"a".as_slice(),
3880                b"b".as_slice(),
3881                b"c".as_slice(),
3882                b"d".as_slice(),
3883            ])
3884        );
3885    }
3886
3887    #[test]
3888    fn cached_schema_plan_does_not_retain_batch_arrays() {
3889        let schema = Arc::new(Schema::new(vec![Field::new(
3890            "a",
3891            ArrowDataType::Int32,
3892            false,
3893        )]));
3894        let array = Arc::new(Int32Array::from(vec![1, 2, 3]));
3895        let batch = RecordBatch::try_new(Arc::clone(&schema), vec![array.clone()]).unwrap();
3896        let mut writer = ArrowWriter::try_new(Vec::new(), schema, None).unwrap();
3897
3898        writer.write(&batch).unwrap();
3899        drop(batch);
3900        assert_eq!(
3901            Arc::strong_count(&array),
3902            1,
3903            "schema and row-group caches must not retain input ArrayRefs"
3904        );
3905        writer.close().unwrap();
3906    }
3907
3908    #[test]
3909    fn arrow_writer_page_size() {
3910        let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)]));
3911
3912        let mut builder = StringBuilder::with_capacity(100, 329 * 10_000);
3913
3914        // Generate an array of 10 unique 10 character string
3915        for i in 0..10 {
3916            let value = i
3917                .to_string()
3918                .repeat(10)
3919                .chars()
3920                .take(10)
3921                .collect::<String>();
3922
3923            builder.append_value(value);
3924        }
3925
3926        let array = Arc::new(builder.finish());
3927
3928        let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
3929
3930        let file = tempfile::tempfile().unwrap();
3931
3932        // Set everything very low so we fallback to PLAIN encoding after the first row
3933        let props = WriterProperties::builder()
3934            .set_data_page_size_limit(1)
3935            .set_dictionary_page_size_limit(1)
3936            .set_write_batch_size(1)
3937            .build();
3938
3939        let mut writer =
3940            ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props))
3941                .expect("Unable to write file");
3942        writer.write(&batch).unwrap();
3943        writer.close().unwrap();
3944
3945        let options = ReadOptionsBuilder::new().with_page_index().build();
3946        let reader =
3947            SerializedFileReader::new_with_options(file.try_clone().unwrap(), options).unwrap();
3948
3949        let column = reader.metadata().row_group(0).columns();
3950
3951        assert_eq!(column.len(), 1);
3952
3953        // We should write one row before falling back to PLAIN encoding so there should still be a
3954        // dictionary page.
3955        assert!(
3956            column[0].dictionary_page_offset().is_some(),
3957            "Expected a dictionary page"
3958        );
3959
3960        let page_index = reader
3961            .metadata()
3962            .page_index()
3963            .expect("page index should be present");
3964        let page_locations = page_index
3965            .page_locations(0, 0)
3966            .expect("page locations should exist");
3967
3968        // We should fallback to PLAIN encoding after the first row and our max page size is 1 bytes
3969        // so we expect one dictionary encoded page and then a page per row thereafter.
3970        assert_eq!(
3971            page_locations.len(),
3972            10,
3973            "Expected 10 pages but got {page_locations:#?}"
3974        );
3975    }
3976
3977    const MEDIUM_SIZE: usize = 63;
3978
3979    fn check_bloom_filter<T: AsBytes>(
3980        files: Vec<Bytes>,
3981        file_column: String,
3982        positive_values: Vec<T>,
3983        negative_values: Vec<T>,
3984    ) {
3985        files.into_iter().take(1).for_each(|file| {
3986            let file_reader = SerializedFileReader::new_with_options(
3987                file,
3988                ReadOptionsBuilder::new()
3989                    .with_reader_properties(
3990                        ReaderProperties::builder()
3991                            .set_read_bloom_filter(true)
3992                            .build(),
3993                    )
3994                    .build(),
3995            )
3996            .expect("Unable to open file as Parquet");
3997            let metadata = file_reader.metadata();
3998
3999            // Gets bloom filters from all row groups.
4000            let mut bloom_filters: Vec<_> = vec![];
4001            for (ri, row_group) in metadata.row_groups().iter().enumerate() {
4002                if let Some((column_index, _)) = row_group
4003                    .columns()
4004                    .iter()
4005                    .enumerate()
4006                    .find(|(_, column)| column.column_path().string() == file_column)
4007                {
4008                    let row_group_reader = file_reader
4009                        .get_row_group(ri)
4010                        .expect("Unable to read row group");
4011                    if let Some(sbbf) = row_group_reader.get_column_bloom_filter(column_index) {
4012                        bloom_filters.push(sbbf.clone());
4013                    } else {
4014                        panic!("No bloom filter for column named {file_column} found");
4015                    }
4016                } else {
4017                    panic!("No column named {file_column} found");
4018                }
4019            }
4020
4021            positive_values.iter().for_each(|value| {
4022                let found = bloom_filters.iter().find(|sbbf| sbbf.check(value));
4023                assert!(
4024                    found.is_some(),
4025                    "{}",
4026                    format!("Value {:?} should be in bloom filter", value.as_bytes())
4027                );
4028            });
4029
4030            negative_values.iter().for_each(|value| {
4031                let found = bloom_filters.iter().find(|sbbf| sbbf.check(value));
4032                assert!(
4033                    found.is_none(),
4034                    "{}",
4035                    format!("Value {:?} should not be in bloom filter", value.as_bytes())
4036                );
4037            });
4038        });
4039    }
4040
4041    #[test]
4042    #[cfg_attr(miri, ignore)] // Takes too long
4043    fn bool_large_single_column() {
4044        let values = Arc::new(
4045            [None, Some(true), Some(false)]
4046                .iter()
4047                .cycle()
4048                .copied()
4049                .take(200_000)
4050                .collect::<BooleanArray>(),
4051        );
4052        let schema = Schema::new(vec![Field::new("col", values.data_type().clone(), true)]);
4053        let expected_batch = RecordBatch::try_new(Arc::new(schema), vec![values]).unwrap();
4054        let file = tempfile::tempfile().unwrap();
4055
4056        let mut writer =
4057            ArrowWriter::try_new(file.try_clone().unwrap(), expected_batch.schema(), None)
4058                .expect("Unable to write file");
4059        writer.write(&expected_batch).unwrap();
4060        writer.close().unwrap();
4061    }
4062
4063    #[test]
4064    fn check_page_offset_index_with_nan() {
4065        let values = Arc::new(Float64Array::from(vec![f64::NAN; 10]));
4066        let schema = Schema::new(vec![Field::new("col", DataType::Float64, true)]);
4067        let batch = RecordBatch::try_new(Arc::new(schema), vec![values]).unwrap();
4068
4069        let mut out = Vec::with_capacity(1024);
4070        let mut writer =
4071            ArrowWriter::try_new(&mut out, batch.schema(), None).expect("Unable to write file");
4072        writer.write(&batch).unwrap();
4073        let file_meta_data = writer.close().unwrap();
4074        for row_group in file_meta_data.row_groups() {
4075            for column in row_group.columns() {
4076                assert!(column.offset_index_offset().is_some());
4077                assert!(column.offset_index_length().is_some());
4078                assert!(column.column_index_offset().is_some());
4079                assert!(column.column_index_length().is_some());
4080            }
4081        }
4082        if let Some(page_index) = file_meta_data.page_index() {
4083            for rg in 0..file_meta_data.num_row_groups() {
4084                for col in 0..file_meta_data.row_group(rg).num_columns() {
4085                    let idx = page_index
4086                        .column_index(rg, col)
4087                        .expect("column index should exist");
4088                    assert!(idx.nan_counts().is_some());
4089                    let ColumnIndexMetaData::DOUBLE(float_idx) = idx else {
4090                        panic!("expected double statistics")
4091                    };
4092                    for i in 0..idx.num_pages() as usize {
4093                        assert_eq!(float_idx.nan_count(i), Some(10));
4094                        assert_eq!(
4095                            f64::NAN.total_cmp(float_idx.min_value(i).unwrap()),
4096                            Ordering::Equal
4097                        );
4098                        assert_eq!(
4099                            f64::NAN.total_cmp(float_idx.max_value(i).unwrap()),
4100                            Ordering::Equal
4101                        );
4102                    }
4103                }
4104            }
4105        } else {
4106            panic!("page index should be present");
4107        }
4108    }
4109
4110    #[test]
4111    fn check_page_offset_index_with_mixed_nan() {
4112        let schema = Arc::new(Schema::new(vec![Field::new(
4113            "col",
4114            DataType::Float64,
4115            true,
4116        )]));
4117
4118        let mut out = Vec::with_capacity(1024);
4119        let props = WriterProperties::builder()
4120            .set_data_page_row_count_limit(10)
4121            .build();
4122        let mut writer = ArrowWriter::try_new(&mut out, schema.clone(), Some(props))
4123            .expect("Unable to write file");
4124
4125        // write a page of all NaN (since batch min and max are NaN, global min/max are NaN)
4126        let values = Arc::new(Float64Array::from(vec![f64::NAN; 10]));
4127        let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
4128        writer.write(&batch).unwrap();
4129
4130        // write a page of all -NaN (batch min/max is -NaN, should update global min to -NaN)
4131        let values = Arc::new(Float64Array::from(vec![-f64::NAN; 10]));
4132        let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
4133        writer.write(&batch).unwrap();
4134
4135        // write a page of all 0 (non-NaN should override global min/max, now 0/0)
4136        let values = Arc::new(Float64Array::from(vec![0_f64; 10]));
4137        let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
4138        writer.write(&batch).unwrap();
4139
4140        // write a mixed page (should now have min -1, max 1)
4141        let values = Arc::new(Float64Array::from(vec![
4142            -1.0,
4143            0.0,
4144            f64::NAN,
4145            -f64::NAN,
4146            1.0,
4147        ]));
4148        let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
4149        writer.write(&batch).unwrap();
4150
4151        let file_meta_data = writer.close().unwrap();
4152
4153        // check the column chunk stats are correct
4154        let col_stats = file_meta_data
4155            .row_group(0)
4156            .column(0)
4157            .statistics()
4158            .expect("missing column chunk statistics");
4159
4160        assert_eq!(col_stats.nan_count_opt(), Some(22));
4161        assert_eq!(col_stats.min_bytes_opt(), Some((-1.0f64).as_bytes()));
4162        assert_eq!(col_stats.max_bytes_opt(), Some(1.0f64.as_bytes()));
4163
4164        assert!(file_meta_data.page_index().is_some());
4165        let col_idx = &file_meta_data.page_index().unwrap().column_index(0, 0);
4166        assert_eq!(col_idx.as_ref().unwrap().num_pages(), 4);
4167
4168        // test each page
4169        let Some(ColumnIndexMetaData::DOUBLE(float_idx)) = col_idx else {
4170            panic!("expected double statistics")
4171        };
4172
4173        assert_eq!(float_idx.nan_counts, Some(vec![10, 10, 0, 2]));
4174        assert_eq!(
4175            f64::NAN.total_cmp(float_idx.min_value(0).unwrap()),
4176            Ordering::Equal
4177        );
4178        assert_eq!(
4179            f64::NAN.total_cmp(float_idx.max_value(0).unwrap()),
4180            Ordering::Equal
4181        );
4182        assert_eq!(
4183            (-f64::NAN).total_cmp(float_idx.min_value(1).unwrap()),
4184            Ordering::Equal
4185        );
4186        assert_eq!(
4187            (-f64::NAN).total_cmp(float_idx.max_value(1).unwrap()),
4188            Ordering::Equal
4189        );
4190        assert_eq!(float_idx.min_value(2), Some(&0.0));
4191        assert_eq!(float_idx.max_value(2), Some(&0.0));
4192        assert_eq!(float_idx.min_value(3), Some(&-1.0));
4193        assert_eq!(float_idx.max_value(3), Some(&1.0));
4194    }
4195
4196    #[test]
4197    #[should_panic(
4198        expected = "Attempting to write an Arrow interval type MonthDayNano to parquet that is not yet implemented"
4199    )]
4200    fn interval_month_day_nano_single_column() {
4201        required_and_optional::<IntervalMonthDayNanoArray, _>(vec![
4202            IntervalMonthDayNano::new(0, 1, 5),
4203            IntervalMonthDayNano::new(0, 3, 2),
4204            IntervalMonthDayNano::new(3, -2, -5),
4205            IntervalMonthDayNano::new(-200, 4, -1),
4206        ]);
4207    }
4208
4209    #[test]
4210    #[cfg_attr(miri, ignore)] // Takes too long
4211    fn i32_column_bloom_filter_at_end() {
4212        let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
4213        let files = RoundTripTest::new(array)
4214            .with_nullable(false)
4215            .with_bloom_filter(true)
4216            .with_bloom_filter_position(BloomFilterPosition::End)
4217            .run();
4218
4219        check_bloom_filter(
4220            files,
4221            "col".to_string(),
4222            (0..SMALL_SIZE as i32).collect(),
4223            (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
4224        );
4225    }
4226
4227    #[test]
4228    #[cfg_attr(miri, ignore)] // Takes too long
4229    fn i32_column_bloom_filter() {
4230        let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
4231        let files = RoundTripTest::new(array)
4232            .with_nullable(false)
4233            .with_bloom_filter(true)
4234            .run();
4235
4236        check_bloom_filter(
4237            files,
4238            "col".to_string(),
4239            (0..SMALL_SIZE as i32).collect(),
4240            (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
4241        );
4242    }
4243
4244    fn write_with_bloom_filter(array: ArrayRef, dictionary_page_size_limit: usize) -> Bytes {
4245        let schema = Arc::new(Schema::new(vec![Field::new(
4246            "col",
4247            array.data_type().clone(),
4248            false,
4249        )]));
4250        let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap();
4251        let props = WriterProperties::builder()
4252            .set_dictionary_enabled(true)
4253            .set_dictionary_page_size_limit(dictionary_page_size_limit)
4254            .set_write_batch_size(256)
4255            .set_bloom_filter_enabled(true)
4256            .build();
4257        let mut buf = Vec::new();
4258        let mut writer = ArrowWriter::try_new(&mut buf, schema, Some(props)).unwrap();
4259        writer.write(&batch).unwrap();
4260        writer.close().unwrap();
4261        Bytes::from(buf)
4262    }
4263
4264    fn data_page_encoding_mask(file: &Bytes) -> EncodingMask {
4265        let metadata = ParquetMetaDataReader::new().parse_and_finish(file).unwrap();
4266        *metadata
4267            .row_group(0)
4268            .column(0)
4269            .page_encoding_stats_mask()
4270            .unwrap()
4271    }
4272
4273    /// While a column is dictionary encoded the bloom filter is populated from the dictionary
4274    /// when it is flushed, so a chunk that stays dictionary encoded must still contain every value.
4275    #[test]
4276    fn string_column_bloom_filter_populated_from_dictionary() {
4277        let values: Vec<String> = (0..2000).map(|i| format!("value-{}", i % 10)).collect();
4278        let array = Arc::new(StringArray::from_iter_values(&values));
4279        let file = write_with_bloom_filter(array, 1024 * 1024);
4280        assert!(data_page_encoding_mask(&file).is_only(Encoding::RLE_DICTIONARY));
4281
4282        check_bloom_filter(
4283            vec![file],
4284            "col".to_string(),
4285            (0..10).map(|i| format!("value-{i}").into_bytes()).collect(),
4286            (10..20)
4287                .map(|i| format!("value-{i}").into_bytes())
4288                .collect(),
4289        );
4290    }
4291
4292    /// After falling back from dictionary encoding the filter holds the dictionary's values
4293    /// and every value written plain afterwards.
4294    #[test]
4295    fn string_column_bloom_filter_across_dictionary_fallback() {
4296        let values: Vec<String> = (0..2000).map(|i| format!("value-{i}")).collect();
4297        let array = Arc::new(StringArray::from_iter_values(&values));
4298        let file = write_with_bloom_filter(array, 1024);
4299        let encodings = data_page_encoding_mask(&file);
4300        assert!(
4301            encodings.is_set(Encoding::RLE_DICTIONARY) && encodings.is_set(Encoding::PLAIN),
4302            "expected dictionary and plain data pages, got {encodings:?}"
4303        );
4304
4305        check_bloom_filter(
4306            vec![file],
4307            "col".to_string(),
4308            values.into_iter().map(String::into_bytes).collect(),
4309            (2000..2010)
4310                .map(|i| format!("value-{i}").into_bytes())
4311                .collect(),
4312        );
4313    }
4314
4315    #[test]
4316    fn i64_column_bloom_filter_populated_from_dictionary() {
4317        let array = Arc::new(Int64Array::from_iter_values((0..2000).map(|i| i % 10)));
4318        let file = write_with_bloom_filter(array, 1024 * 1024);
4319        assert!(data_page_encoding_mask(&file).is_only(Encoding::RLE_DICTIONARY));
4320
4321        check_bloom_filter(
4322            vec![file],
4323            "col".to_string(),
4324            (0..10i64).collect(),
4325            (10..20i64).collect(),
4326        );
4327    }
4328
4329    #[test]
4330    fn i64_column_bloom_filter_across_dictionary_fallback() {
4331        let array = Arc::new(Int64Array::from_iter_values(0..2000i64));
4332        let file = write_with_bloom_filter(array, 1024);
4333        let encodings = data_page_encoding_mask(&file);
4334        assert!(
4335            encodings.is_set(Encoding::RLE_DICTIONARY) && encodings.is_set(Encoding::PLAIN),
4336            "expected dictionary and plain data pages, got {encodings:?}"
4337        );
4338
4339        check_bloom_filter(
4340            vec![file],
4341            "col".to_string(),
4342            (0..2000i64).collect(),
4343            (2000..2010i64).collect(),
4344        );
4345    }
4346
4347    /// Test that bloom filter folding produces correct results even when
4348    /// the configured NDV differs significantly from actual NDV.
4349    /// A large NDV means a larger initial filter that gets folded down;
4350    /// a small NDV means a smaller initial filter.
4351    #[test]
4352    #[cfg_attr(miri, ignore)] // Takes too long
4353    fn i32_column_bloom_filter_fixed_ndv() {
4354        let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
4355
4356        // NDV much larger than actual distinct values — tests folding a large filter down
4357        let files = RoundTripTest::new(array.clone())
4358            .with_nullable(false)
4359            .with_bloom_filter(true)
4360            .with_bloom_filter_ndv(1_000_000)
4361            .run();
4362
4363        check_bloom_filter(
4364            files,
4365            "col".to_string(),
4366            (0..SMALL_SIZE as i32).collect(),
4367            (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
4368        );
4369
4370        // NDV smaller than actual distinct values — tests the underestimate path
4371        let files = RoundTripTest::new(array)
4372            .with_nullable(false)
4373            .with_bloom_filter(true)
4374            .with_bloom_filter_ndv(3)
4375            .run();
4376
4377        check_bloom_filter(
4378            files,
4379            "col".to_string(),
4380            (0..SMALL_SIZE as i32).collect(),
4381            (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
4382        );
4383    }
4384
4385    #[test]
4386    #[cfg_attr(miri, ignore)] // Takes too long
4387    fn binary_column_bloom_filter() {
4388        let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
4389        let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
4390        let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
4391
4392        let array = Arc::new(BinaryArray::from_iter_values(many_vecs_iter));
4393        let files = RoundTripTest::new(array)
4394            .with_nullable(false)
4395            .with_bloom_filter(true)
4396            .run();
4397
4398        check_bloom_filter(
4399            files,
4400            "col".to_string(),
4401            many_vecs,
4402            vec![vec![(SMALL_SIZE + 1) as u8]],
4403        );
4404    }
4405
4406    #[test]
4407    #[cfg_attr(miri, ignore)] // Takes too long
4408    fn empty_string_null_column_bloom_filter() {
4409        let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
4410        let raw_strs = raw_values.iter().map(|s| s.as_str());
4411
4412        let array = Arc::new(StringArray::from_iter_values(raw_strs));
4413        let files = RoundTripTest::new(array)
4414            .with_nullable(false)
4415            .with_bloom_filter(true)
4416            .run();
4417
4418        let optional_raw_values: Vec<_> = raw_values
4419            .iter()
4420            .enumerate()
4421            .filter_map(|(i, v)| if i % 2 == 0 { None } else { Some(v.as_str()) })
4422            .collect();
4423        // For null slots, empty string should not be in bloom filter.
4424        check_bloom_filter(files, "col".to_string(), optional_raw_values, vec![""]);
4425    }
4426
4427    #[test]
4428    fn list_and_map_coerced_names() {
4429        // Create map and list with non-Parquet naming
4430        let list_field =
4431            Field::new_list("my_list", Field::new("item", DataType::Int32, false), false);
4432        let map_field = Field::new_map(
4433            "my_map",
4434            "my_entries",
4435            Field::new("my_keys", DataType::Int32, false),
4436            Field::new("my_values", DataType::Int32, true),
4437            false,
4438            true,
4439        );
4440
4441        let list_array = create_random_array(&list_field, 100, 0.0, 0.0).unwrap();
4442        let map_array = create_random_array(&map_field, 100, 0.0, 0.0).unwrap();
4443
4444        let arrow_schema = Arc::new(Schema::new(vec![list_field, map_field]));
4445
4446        // Write data to Parquet but coerce names to match spec
4447        let props = Some(WriterProperties::builder().set_coerce_types(true).build());
4448        let file = tempfile::tempfile().unwrap();
4449        let mut writer =
4450            ArrowWriter::try_new(file.try_clone().unwrap(), arrow_schema.clone(), props).unwrap();
4451
4452        let batch = RecordBatch::try_new(arrow_schema, vec![list_array, map_array]).unwrap();
4453        writer.write(&batch).unwrap();
4454        let file_metadata = writer.close().unwrap();
4455
4456        let schema = file_metadata.file_metadata().schema();
4457        // Coerced name of "item" should be "element"
4458        let list_field = &schema.get_fields()[0].get_fields()[0];
4459        assert_eq!(list_field.get_fields()[0].name(), "element");
4460
4461        let map_field = &schema.get_fields()[1].get_fields()[0];
4462        // Coerced name of "entries" should be "key_value"
4463        assert_eq!(map_field.name(), "key_value");
4464        // Coerced name of "my_keys" should be "key"
4465        assert_eq!(map_field.get_fields()[0].name(), "key");
4466        // Coerced name of "my_values" should be "value"
4467        assert_eq!(map_field.get_fields()[1].name(), "value");
4468
4469        // Double check schema after reading from the file
4470        let reader = SerializedFileReader::new(file).unwrap();
4471        let file_schema = reader.metadata().file_metadata().schema();
4472        let fields = file_schema.get_fields();
4473        let list_field = &fields[0].get_fields()[0];
4474        assert_eq!(list_field.get_fields()[0].name(), "element");
4475        let map_field = &fields[1].get_fields()[0];
4476        assert_eq!(map_field.name(), "key_value");
4477        assert_eq!(map_field.get_fields()[0].name(), "key");
4478        assert_eq!(map_field.get_fields()[1].name(), "value");
4479    }
4480
4481    #[test]
4482    #[cfg_attr(miri, ignore)] // Takes too long
4483    fn fallback_flush_data_page() {
4484        //tests if the Fallback::flush_data_page clears all buffers correctly
4485        let raw_values: Vec<_> = (0..MEDIUM_SIZE).map(|i| i.to_string()).collect();
4486        let values = Arc::new(StringArray::from(raw_values));
4487        let encodings = vec![
4488            Encoding::DELTA_BYTE_ARRAY,
4489            Encoding::DELTA_LENGTH_BYTE_ARRAY,
4490        ];
4491        let data_type = values.data_type().clone();
4492        let schema = Arc::new(Schema::new(vec![Field::new("col", data_type, false)]));
4493        let expected_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4494
4495        let row_group_sizes = [1024, SMALL_SIZE, SMALL_SIZE / 2, SMALL_SIZE / 2 + 1, 10];
4496        let data_page_size_limit: usize = 32;
4497        let write_batch_size: usize = 16;
4498
4499        for encoding in &encodings {
4500            for row_group_size in row_group_sizes {
4501                let props = WriterProperties::builder()
4502                    .set_writer_version(WriterVersion::PARQUET_2_0)
4503                    .set_max_row_group_row_count(Some(row_group_size))
4504                    .set_dictionary_enabled(false)
4505                    .set_encoding(*encoding)
4506                    .set_data_page_size_limit(data_page_size_limit)
4507                    .set_write_batch_size(write_batch_size)
4508                    .build();
4509
4510                roundtrip_opts_with_array_validation(&expected_batch, props, |a, b| {
4511                    let string_array_a = StringArray::from(a.clone());
4512                    let string_array_b = StringArray::from(b.clone());
4513                    let vec_a: Vec<&str> = string_array_a.iter().map(|v| v.unwrap()).collect();
4514                    let vec_b: Vec<&str> = string_array_b.iter().map(|v| v.unwrap()).collect();
4515                    assert_eq!(
4516                        vec_a, vec_b,
4517                        "failed for encoder: {encoding:?} and row_group_size: {row_group_size:?}"
4518                    );
4519                });
4520            }
4521        }
4522    }
4523
4524    #[test]
4525    #[cfg_attr(miri, ignore)] // Takes too long
4526    fn u32_min_max() {
4527        // check values roundtrip through parquet
4528        let src = [
4529            u32::MIN,
4530            1,
4531            (i32::MAX as u32) - 1,
4532            i32::MAX as u32,
4533            (i32::MAX as u32) + 1,
4534            u32::MAX - 1,
4535            u32::MAX,
4536        ];
4537        let values = Arc::new(UInt32Array::from_iter_values(src.iter().copied()));
4538        let files = RoundTripTest::new(values).with_nullable(false).run();
4539
4540        for file in files {
4541            // check statistics are valid
4542            let reader = SerializedFileReader::new(file).unwrap();
4543            let metadata = reader.metadata();
4544
4545            let mut row_offset = 0;
4546            for row_group in metadata.row_groups() {
4547                assert_eq!(row_group.num_columns(), 1);
4548                let column = row_group.column(0);
4549
4550                let num_values = column.num_values() as usize;
4551                let src_slice = &src[row_offset..row_offset + num_values];
4552                row_offset += column.num_values() as usize;
4553
4554                let stats = column.statistics().unwrap();
4555                if let Statistics::Int32(stats) = stats {
4556                    assert_eq!(
4557                        *stats.min_opt().unwrap() as u32,
4558                        *src_slice.iter().min().unwrap()
4559                    );
4560                    assert_eq!(
4561                        *stats.max_opt().unwrap() as u32,
4562                        *src_slice.iter().max().unwrap()
4563                    );
4564                } else {
4565                    panic!("Statistics::Int32 missing")
4566                }
4567            }
4568        }
4569    }
4570
4571    #[test]
4572    #[cfg_attr(miri, ignore)] // Takes too long
4573    fn u64_min_max() {
4574        // check values roundtrip through parquet
4575        let src = [
4576            u64::MIN,
4577            1,
4578            (i64::MAX as u64) - 1,
4579            i64::MAX as u64,
4580            (i64::MAX as u64) + 1,
4581            u64::MAX - 1,
4582            u64::MAX,
4583        ];
4584        let values = Arc::new(UInt64Array::from_iter_values(src.iter().copied()));
4585        let files = RoundTripTest::new(values).with_nullable(false).run();
4586
4587        for file in files {
4588            // check statistics are valid
4589            let reader = SerializedFileReader::new(file).unwrap();
4590            let metadata = reader.metadata();
4591
4592            let mut row_offset = 0;
4593            for row_group in metadata.row_groups() {
4594                assert_eq!(row_group.num_columns(), 1);
4595                let column = row_group.column(0);
4596
4597                let num_values = column.num_values() as usize;
4598                let src_slice = &src[row_offset..row_offset + num_values];
4599                row_offset += column.num_values() as usize;
4600
4601                let stats = column.statistics().unwrap();
4602                if let Statistics::Int64(stats) = stats {
4603                    assert_eq!(
4604                        *stats.min_opt().unwrap() as u64,
4605                        *src_slice.iter().min().unwrap()
4606                    );
4607                    assert_eq!(
4608                        *stats.max_opt().unwrap() as u64,
4609                        *src_slice.iter().max().unwrap()
4610                    );
4611                } else {
4612                    panic!("Statistics::Int64 missing")
4613                }
4614            }
4615        }
4616    }
4617
4618    #[test]
4619    #[cfg_attr(miri, ignore)] // Takes too long
4620    fn statistics_null_counts_only_nulls() {
4621        // check that null-count statistics for "only NULL"-columns are correct
4622        let values = Arc::new(UInt64Array::from(vec![None, None]));
4623        let files = RoundTripTest::new(values).run();
4624
4625        for file in files {
4626            // check statistics are valid
4627            let reader = SerializedFileReader::new(file).unwrap();
4628            let metadata = reader.metadata();
4629            assert_eq!(metadata.num_row_groups(), 1);
4630            let row_group = metadata.row_group(0);
4631            assert_eq!(row_group.num_columns(), 1);
4632            let column = row_group.column(0);
4633            let stats = column.statistics().unwrap();
4634            assert_eq!(stats.null_count_opt(), Some(2));
4635        }
4636    }
4637
4638    fn row_group_sizes(metadata: &ParquetMetaData) -> Vec<i64> {
4639        metadata.row_groups().iter().map(|x| x.num_rows()).collect()
4640    }
4641
4642    #[test]
4643    fn test_aggregates_records() {
4644        let arrays = [
4645            Int32Array::from((0..100).collect::<Vec<_>>()),
4646            Int32Array::from((0..50).collect::<Vec<_>>()),
4647            Int32Array::from((200..500).collect::<Vec<_>>()),
4648        ];
4649
4650        let schema = Arc::new(Schema::new(vec![Field::new(
4651            "int",
4652            ArrowDataType::Int32,
4653            false,
4654        )]));
4655
4656        let file = tempfile::tempfile().unwrap();
4657
4658        let props = WriterProperties::builder()
4659            .set_max_row_group_row_count(Some(200))
4660            .build();
4661
4662        let mut writer =
4663            ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
4664
4665        for array in arrays {
4666            let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
4667            writer.write(&batch).unwrap();
4668        }
4669
4670        writer.close().unwrap();
4671
4672        let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
4673        assert_eq!(&row_group_sizes(builder.metadata()), &[200, 200, 50]);
4674
4675        let batches = builder
4676            .with_batch_size(100)
4677            .build()
4678            .unwrap()
4679            .collect::<ArrowResult<Vec<_>>>()
4680            .unwrap();
4681
4682        assert_eq!(batches.len(), 5);
4683        assert!(batches.iter().all(|x| x.num_columns() == 1));
4684
4685        let batch_sizes: Vec<_> = batches.iter().map(|x| x.num_rows()).collect();
4686
4687        assert_eq!(&batch_sizes, &[100, 100, 100, 100, 50]);
4688
4689        let values: Vec<_> = batches
4690            .iter()
4691            .flat_map(|x| {
4692                x.column(0)
4693                    .as_any()
4694                    .downcast_ref::<Int32Array>()
4695                    .unwrap()
4696                    .values()
4697                    .iter()
4698                    .copied()
4699            })
4700            .collect();
4701
4702        let expected_values: Vec<_> = [0..100, 0..50, 200..500].into_iter().flatten().collect();
4703        assert_eq!(&values, &expected_values)
4704    }
4705
4706    #[test]
4707    fn complex_aggregate() {
4708        // Tests aggregating nested data
4709        let field_a = Arc::new(Field::new("leaf_a", DataType::Int32, false));
4710        let field_b = Arc::new(Field::new("leaf_b", DataType::Int32, true));
4711        let struct_a = Arc::new(Field::new(
4712            "struct_a",
4713            DataType::Struct(vec![field_a.clone(), field_b.clone()].into()),
4714            true,
4715        ));
4716
4717        let list_a = Arc::new(Field::new("list", DataType::List(struct_a), true));
4718        let struct_b = Arc::new(Field::new(
4719            "struct_b",
4720            DataType::Struct(vec![list_a.clone()].into()),
4721            false,
4722        ));
4723
4724        let schema = Arc::new(Schema::new(vec![struct_b]));
4725
4726        // create nested data
4727        let field_a_array = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
4728        let field_b_array =
4729            Int32Array::from_iter(vec![Some(1), None, Some(2), None, None, Some(6)]);
4730
4731        let struct_a_array = StructArray::from(vec![
4732            (field_a.clone(), Arc::new(field_a_array) as ArrayRef),
4733            (field_b.clone(), Arc::new(field_b_array) as ArrayRef),
4734        ]);
4735
4736        let list_data = ArrayDataBuilder::new(list_a.data_type().clone())
4737            .len(5)
4738            .add_buffer(Buffer::from_iter(vec![
4739                0_i32, 1_i32, 1_i32, 3_i32, 3_i32, 5_i32,
4740            ]))
4741            .null_bit_buffer(Some(Buffer::from_iter(vec![
4742                true, false, true, false, true,
4743            ])))
4744            .child_data(vec![struct_a_array.into_data()])
4745            .build()
4746            .unwrap();
4747
4748        let list_a_array = Arc::new(ListArray::from(list_data)) as ArrayRef;
4749        let struct_b_array = StructArray::from(vec![(list_a.clone(), list_a_array)]);
4750
4751        let batch1 =
4752            RecordBatch::try_from_iter(vec![("struct_b", Arc::new(struct_b_array) as ArrayRef)])
4753                .unwrap();
4754
4755        let field_a_array = Int32Array::from(vec![6, 7, 8, 9, 10]);
4756        let field_b_array = Int32Array::from_iter(vec![None, None, None, Some(1), None]);
4757
4758        let struct_a_array = StructArray::from(vec![
4759            (field_a, Arc::new(field_a_array) as ArrayRef),
4760            (field_b, Arc::new(field_b_array) as ArrayRef),
4761        ]);
4762
4763        let list_data = ArrayDataBuilder::new(list_a.data_type().clone())
4764            .len(2)
4765            .add_buffer(Buffer::from_iter(vec![0_i32, 4_i32, 5_i32]))
4766            .child_data(vec![struct_a_array.into_data()])
4767            .build()
4768            .unwrap();
4769
4770        let list_a_array = Arc::new(ListArray::from(list_data)) as ArrayRef;
4771        let struct_b_array = StructArray::from(vec![(list_a, list_a_array)]);
4772
4773        let batch2 =
4774            RecordBatch::try_from_iter(vec![("struct_b", Arc::new(struct_b_array) as ArrayRef)])
4775                .unwrap();
4776
4777        let batches = &[batch1, batch2];
4778
4779        // Verify data is as expected
4780
4781        let expected = r"
4782            +-------------------------------------------------------------------------------------------------------+
4783            | struct_b                                                                                              |
4784            +-------------------------------------------------------------------------------------------------------+
4785            | {list: [{leaf_a: 1, leaf_b: 1}]}                                                                      |
4786            | {list: }                                                                                              |
4787            | {list: [{leaf_a: 2, leaf_b: }, {leaf_a: 3, leaf_b: 2}]}                                               |
4788            | {list: }                                                                                              |
4789            | {list: [{leaf_a: 4, leaf_b: }, {leaf_a: 5, leaf_b: }]}                                                |
4790            | {list: [{leaf_a: 6, leaf_b: }, {leaf_a: 7, leaf_b: }, {leaf_a: 8, leaf_b: }, {leaf_a: 9, leaf_b: 1}]} |
4791            | {list: [{leaf_a: 10, leaf_b: }]}                                                                      |
4792            +-------------------------------------------------------------------------------------------------------+
4793        ".trim().split('\n').map(|x| x.trim()).collect::<Vec<_>>().join("\n");
4794
4795        let actual = pretty_format_batches(batches).unwrap().to_string();
4796        assert_eq!(actual, expected);
4797
4798        // Write data
4799        let file = tempfile::tempfile().unwrap();
4800        let props = WriterProperties::builder()
4801            .set_max_row_group_row_count(Some(6))
4802            .build();
4803
4804        let mut writer =
4805            ArrowWriter::try_new(file.try_clone().unwrap(), schema, Some(props)).unwrap();
4806
4807        for batch in batches {
4808            writer.write(batch).unwrap();
4809        }
4810        writer.close().unwrap();
4811
4812        // Read Data
4813        // Should have written entire first batch and first row of second to the first row group
4814        // leaving a single row in the second row group
4815
4816        let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
4817        assert_eq!(&row_group_sizes(builder.metadata()), &[6, 1]);
4818
4819        let batches = builder
4820            .with_batch_size(2)
4821            .build()
4822            .unwrap()
4823            .collect::<ArrowResult<Vec<_>>>()
4824            .unwrap();
4825
4826        assert_eq!(batches.len(), 4);
4827        let batch_counts: Vec<_> = batches.iter().map(|x| x.num_rows()).collect();
4828        assert_eq!(&batch_counts, &[2, 2, 2, 1]);
4829
4830        let actual = pretty_format_batches(&batches).unwrap().to_string();
4831        assert_eq!(actual, expected);
4832    }
4833
4834    #[test]
4835    fn test_arrow_writer_metadata() {
4836        let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
4837        let file_schema = batch_schema.clone().with_metadata([("foo", "bar")]);
4838
4839        let batch = RecordBatch::try_new(
4840            Arc::new(batch_schema),
4841            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
4842        )
4843        .unwrap();
4844
4845        let mut buf = Vec::with_capacity(1024);
4846        let mut writer = ArrowWriter::try_new(&mut buf, Arc::new(file_schema), None).unwrap();
4847        writer.write(&batch).unwrap();
4848        writer.close().unwrap();
4849    }
4850
4851    #[test]
4852    fn in_progress_accounting() {
4853        // define schema
4854        let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
4855
4856        // create some data
4857        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
4858
4859        // build a record batch
4860        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
4861
4862        let mut writer = ArrowWriter::try_new(vec![], batch.schema(), None).unwrap();
4863
4864        // starts empty
4865        assert_eq!(writer.in_progress_size(), 0);
4866        assert_eq!(writer.in_progress_rows(), 0);
4867        assert_eq!(writer.memory_size(), 0);
4868        assert_eq!(writer.bytes_written(), 4); // Initial header
4869        writer.write(&batch).unwrap();
4870
4871        // updated on write
4872        let initial_size = writer.in_progress_size();
4873        assert!(initial_size > 0);
4874        assert_eq!(writer.in_progress_rows(), 5);
4875        let initial_memory = writer.memory_size();
4876        assert!(initial_memory > 0);
4877        // memory estimate is larger than estimated encoded size
4878        assert!(
4879            initial_size <= initial_memory,
4880            "{initial_size} <= {initial_memory}"
4881        );
4882
4883        // updated on second write
4884        writer.write(&batch).unwrap();
4885        assert!(writer.in_progress_size() > initial_size);
4886        assert_eq!(writer.in_progress_rows(), 10);
4887        assert!(writer.memory_size() > initial_memory);
4888        assert!(
4889            writer.in_progress_size() <= writer.memory_size(),
4890            "in_progress_size {} <= memory_size {}",
4891            writer.in_progress_size(),
4892            writer.memory_size()
4893        );
4894
4895        // in progress tracking is cleared, but the overall data written is updated
4896        let pre_flush_bytes_written = writer.bytes_written();
4897        writer.flush().unwrap();
4898        assert_eq!(writer.in_progress_size(), 0);
4899        assert_eq!(writer.memory_size(), 0);
4900        assert!(writer.bytes_written() > pre_flush_bytes_written);
4901
4902        writer.close().unwrap();
4903    }
4904
4905    #[test]
4906    fn test_writer_all_null() {
4907        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
4908        let b = Int32Array::new(vec![0; 5].into(), Some(NullBuffer::new_null(5)));
4909        let batch = RecordBatch::try_from_iter(vec![
4910            ("a", Arc::new(a) as ArrayRef),
4911            ("b", Arc::new(b) as ArrayRef),
4912        ])
4913        .unwrap();
4914
4915        let mut buf = Vec::with_capacity(1024);
4916        let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), None).unwrap();
4917        writer.write(&batch).unwrap();
4918        writer.close().unwrap();
4919
4920        let bytes = Bytes::from(buf);
4921        let options = ReadOptionsBuilder::new().with_page_index().build();
4922        let reader = SerializedFileReader::new_with_options(bytes, options).unwrap();
4923        let index = reader.metadata().page_index().unwrap();
4924
4925        assert_eq!(index.num_data_pages(0, 0), Some(1)); // 1 page
4926        assert_eq!(index.num_data_pages(0, 1), Some(1)); // 1 page
4927    }
4928
4929    #[test]
4930    fn test_disabled_statistics_with_page() {
4931        let file_schema = Schema::new(vec![
4932            Field::new("a", DataType::Utf8, true),
4933            Field::new("b", DataType::Utf8, true),
4934        ]);
4935        let file_schema = Arc::new(file_schema);
4936
4937        let batch = RecordBatch::try_new(
4938            file_schema.clone(),
4939            vec![
4940                Arc::new(StringArray::from(vec!["a", "b", "c", "d"])) as _,
4941                Arc::new(StringArray::from(vec!["w", "x", "y", "z"])) as _,
4942            ],
4943        )
4944        .unwrap();
4945
4946        let props = WriterProperties::builder()
4947            .set_statistics_enabled(EnabledStatistics::None)
4948            .set_column_statistics_enabled("a".into(), EnabledStatistics::Page)
4949            .build();
4950
4951        let mut buf = Vec::with_capacity(1024);
4952        let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), Some(props)).unwrap();
4953        writer.write(&batch).unwrap();
4954
4955        let metadata = writer.close().unwrap();
4956        assert_eq!(metadata.num_row_groups(), 1);
4957        let row_group = metadata.row_group(0);
4958        assert_eq!(row_group.num_columns(), 2);
4959        // Column "a" has both offset and column index, as requested
4960        assert!(row_group.column(0).offset_index_offset().is_some());
4961        assert!(row_group.column(0).column_index_offset().is_some());
4962        // Column "b" should only have offset index
4963        assert!(row_group.column(1).offset_index_offset().is_some());
4964        assert!(row_group.column(1).column_index_offset().is_none());
4965
4966        let options = ReadOptionsBuilder::new().with_page_index().build();
4967        let reader = SerializedFileReader::new_with_options(Bytes::from(buf), options).unwrap();
4968
4969        let row_group = reader.get_row_group(0).unwrap();
4970        let a_col = row_group.metadata().column(0);
4971        let b_col = row_group.metadata().column(1);
4972
4973        // Column chunk of column "a" should have chunk level statistics
4974        if let Statistics::ByteArray(byte_array_stats) = a_col.statistics().unwrap() {
4975            let min = byte_array_stats.min_opt().unwrap();
4976            let max = byte_array_stats.max_opt().unwrap();
4977
4978            assert_eq!(min.as_bytes(), b"a");
4979            assert_eq!(max.as_bytes(), b"d");
4980        } else {
4981            panic!("expecting Statistics::ByteArray");
4982        }
4983
4984        // The column chunk for column "b" shouldn't have statistics
4985        assert!(b_col.statistics().is_none());
4986
4987        let page_index = reader.metadata().page_index().unwrap();
4988
4989        let a_idx = page_index.column_index(0, 0);
4990        assert!(
4991            matches!(a_idx, Some(ColumnIndexMetaData::BYTE_ARRAY(_))),
4992            "{a_idx:?}"
4993        );
4994        let b_idx = page_index.column_index(0, 1);
4995        assert!(b_idx.is_none(), "{b_idx:?}");
4996    }
4997
4998    #[test]
4999    fn test_disabled_statistics_with_chunk() {
5000        let file_schema = Schema::new(vec![
5001            Field::new("a", DataType::Utf8, true),
5002            Field::new("b", DataType::Utf8, true),
5003        ]);
5004        let file_schema = Arc::new(file_schema);
5005
5006        let batch = RecordBatch::try_new(
5007            file_schema.clone(),
5008            vec![
5009                Arc::new(StringArray::from(vec!["a", "b", "c", "d"])) as _,
5010                Arc::new(StringArray::from(vec!["w", "x", "y", "z"])) as _,
5011            ],
5012        )
5013        .unwrap();
5014
5015        let props = WriterProperties::builder()
5016            .set_statistics_enabled(EnabledStatistics::None)
5017            .set_column_statistics_enabled("a".into(), EnabledStatistics::Chunk)
5018            .build();
5019
5020        let mut buf = Vec::with_capacity(1024);
5021        let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), Some(props)).unwrap();
5022        writer.write(&batch).unwrap();
5023
5024        let metadata = writer.close().unwrap();
5025        assert_eq!(metadata.num_row_groups(), 1);
5026        let row_group = metadata.row_group(0);
5027        assert_eq!(row_group.num_columns(), 2);
5028        // Column "a" should only have offset index
5029        assert!(row_group.column(0).offset_index_offset().is_some());
5030        assert!(row_group.column(0).column_index_offset().is_none());
5031        // Column "b" should only have offset index
5032        assert!(row_group.column(1).offset_index_offset().is_some());
5033        assert!(row_group.column(1).column_index_offset().is_none());
5034
5035        let options = ReadOptionsBuilder::new().with_page_index().build();
5036        let reader = SerializedFileReader::new_with_options(Bytes::from(buf), options).unwrap();
5037
5038        let row_group = reader.get_row_group(0).unwrap();
5039        let a_col = row_group.metadata().column(0);
5040        let b_col = row_group.metadata().column(1);
5041
5042        // Column chunk of column "a" should have chunk level statistics
5043        if let Statistics::ByteArray(byte_array_stats) = a_col.statistics().unwrap() {
5044            let min = byte_array_stats.min_opt().unwrap();
5045            let max = byte_array_stats.max_opt().unwrap();
5046
5047            assert_eq!(min.as_bytes(), b"a");
5048            assert_eq!(max.as_bytes(), b"d");
5049        } else {
5050            panic!("expecting Statistics::ByteArray");
5051        }
5052
5053        // The column chunk for column "b"  shouldn't have statistics
5054        assert!(b_col.statistics().is_none());
5055
5056        let page_index = reader.metadata().page_index().unwrap();
5057
5058        let a_idx = page_index.column_index(0, 0);
5059        assert!(a_idx.is_none(), "{a_idx:?}");
5060        let b_idx = page_index.column_index(0, 1);
5061        assert!(b_idx.is_none(), "{b_idx:?}");
5062    }
5063
5064    #[test]
5065    fn test_arrow_writer_skip_metadata() {
5066        let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
5067        let file_schema = Arc::new(batch_schema.clone());
5068
5069        let batch = RecordBatch::try_new(
5070            Arc::new(batch_schema),
5071            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5072        )
5073        .unwrap();
5074        let skip_options = ArrowWriterOptions::new().with_skip_arrow_metadata(true);
5075
5076        let mut buf = Vec::with_capacity(1024);
5077        let mut writer =
5078            ArrowWriter::try_new_with_options(&mut buf, file_schema.clone(), skip_options).unwrap();
5079        writer.write(&batch).unwrap();
5080        writer.close().unwrap();
5081
5082        let bytes = Bytes::from(buf);
5083        let reader_builder = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap();
5084        assert_eq!(file_schema, *reader_builder.schema());
5085        if let Some(key_value_metadata) = reader_builder
5086            .metadata()
5087            .file_metadata()
5088            .key_value_metadata()
5089        {
5090            assert!(
5091                !key_value_metadata
5092                    .iter()
5093                    .any(|kv| kv.key.as_str() == ARROW_SCHEMA_META_KEY)
5094            );
5095        }
5096    }
5097
5098    #[test]
5099    fn test_arrow_writer_skip_path_in_schema() {
5100        let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
5101        let file_schema = Arc::new(batch_schema.clone());
5102
5103        let batch = RecordBatch::try_new(
5104            Arc::new(batch_schema),
5105            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5106        )
5107        .unwrap();
5108
5109        // default options should still write path_in_schema
5110        let skip_options = ArrowWriterOptions::new();
5111
5112        let mut buf = Vec::with_capacity(1024);
5113        let mut writer =
5114            ArrowWriter::try_new_with_options(&mut buf, file_schema.clone(), skip_options).unwrap();
5115        writer.write(&batch).unwrap();
5116        writer.close().unwrap();
5117
5118        // override to not write path_in_schema
5119        let skip_options = ArrowWriterOptions::new().with_properties(
5120            WriterProperties::builder()
5121                .set_write_path_in_schema(false)
5122                .build(),
5123        );
5124
5125        let mut buf2 = Vec::with_capacity(1024);
5126        let mut writer =
5127            ArrowWriter::try_new_with_options(&mut buf2, file_schema.clone(), skip_options)
5128                .unwrap();
5129        writer.write(&batch).unwrap();
5130        writer.close().unwrap();
5131
5132        // buf2 should be a bit smaller due to lack of path_in_schema
5133        assert!(buf.len() > buf2.len());
5134    }
5135
5136    #[test]
5137    fn mismatched_schemas() {
5138        let batch_schema = Schema::new(vec![Field::new("count", DataType::Int32, false)]);
5139        let file_schema = Arc::new(Schema::new(vec![Field::new(
5140            "temperature",
5141            DataType::Float64,
5142            false,
5143        )]));
5144
5145        let batch = RecordBatch::try_new(
5146            Arc::new(batch_schema),
5147            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5148        )
5149        .unwrap();
5150
5151        let mut buf = Vec::with_capacity(1024);
5152        let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), None).unwrap();
5153
5154        let err = writer.write(&batch).unwrap_err().to_string();
5155        assert_eq!(
5156            err,
5157            "Arrow: Incompatible type. Field 'temperature' has type Float64, array has type Int32"
5158        );
5159    }
5160
5161    #[test]
5162    fn test_page_stats_not_written_by_default() {
5163        let string_field = Field::new("a", DataType::Utf8, false);
5164        let schema = Schema::new(vec![string_field]);
5165        let raw_string_values = vec!["Blart Versenwald III"];
5166        let string_values = StringArray::from(raw_string_values.clone());
5167        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(string_values)]).unwrap();
5168
5169        let props = WriterProperties::builder()
5170            .set_statistics_enabled(EnabledStatistics::Page)
5171            .set_dictionary_enabled(false)
5172            .set_encoding(Encoding::PLAIN)
5173            .set_compression(crate::basic::Compression::UNCOMPRESSED)
5174            .build();
5175
5176        let file = roundtrip_opts(&batch, props);
5177
5178        // read file and decode page headers
5179        // Note: use the thrift API as there is no Rust API to access the statistics in the page headers
5180
5181        // decode first page header
5182        let first_page = &file[4..];
5183        let mut prot = ThriftSliceInputProtocol::new(first_page);
5184        let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5185        let stats = hdr.data_page_header.unwrap().statistics;
5186
5187        assert!(stats.is_none());
5188    }
5189
5190    #[test]
5191    fn test_page_stats_when_enabled() {
5192        let string_field = Field::new("a", DataType::Utf8, false);
5193        let schema = Schema::new(vec![string_field]);
5194        let raw_string_values = vec!["Blart Versenwald III", "Andrew Lamb"];
5195        let string_values = StringArray::from(raw_string_values.clone());
5196        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(string_values)]).unwrap();
5197
5198        let props = WriterProperties::builder()
5199            .set_statistics_enabled(EnabledStatistics::Page)
5200            .set_dictionary_enabled(false)
5201            .set_encoding(Encoding::PLAIN)
5202            .set_write_page_header_statistics(true)
5203            .set_compression(crate::basic::Compression::UNCOMPRESSED)
5204            .build();
5205
5206        let file = roundtrip_opts(&batch, props);
5207
5208        // read file and decode page headers
5209        // Note: use the thrift API as there is no Rust API to access the statistics in the page headers
5210
5211        // decode first page header
5212        let first_page = &file[4..];
5213        let mut prot = ThriftSliceInputProtocol::new(first_page);
5214        let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5215        let stats = hdr.data_page_header.unwrap().statistics;
5216
5217        let stats = stats.unwrap();
5218        // check that min/max were actually written to the page
5219        assert!(stats.is_max_value_exact.unwrap());
5220        assert!(stats.is_min_value_exact.unwrap());
5221        assert_eq!(stats.max_value.unwrap(), b"Blart Versenwald III");
5222        assert_eq!(stats.min_value.unwrap(), b"Andrew Lamb");
5223    }
5224
5225    #[test]
5226    fn test_page_stats_truncation() {
5227        let string_field = Field::new("a", DataType::Utf8, false);
5228        let binary_field = Field::new("b", DataType::Binary, false);
5229        let schema = Schema::new(vec![string_field, binary_field]);
5230
5231        let raw_string_values = vec!["Blart Versenwald III"];
5232        let raw_binary_values = [b"Blart Versenwald III".to_vec()];
5233        let raw_binary_value_refs = raw_binary_values
5234            .iter()
5235            .map(|x| x.as_slice())
5236            .collect::<Vec<_>>();
5237
5238        let string_values = StringArray::from(raw_string_values.clone());
5239        let binary_values = BinaryArray::from(raw_binary_value_refs);
5240        let batch = RecordBatch::try_new(
5241            Arc::new(schema),
5242            vec![Arc::new(string_values), Arc::new(binary_values)],
5243        )
5244        .unwrap();
5245
5246        let props = WriterProperties::builder()
5247            .set_statistics_truncate_length(Some(2))
5248            .set_dictionary_enabled(false)
5249            .set_encoding(Encoding::PLAIN)
5250            .set_write_page_header_statistics(true)
5251            .set_compression(crate::basic::Compression::UNCOMPRESSED)
5252            .build();
5253
5254        let file = roundtrip_opts(&batch, props);
5255
5256        // read file and decode page headers
5257        // Note: use the thrift API as there is no Rust API to access the statistics in the page headers
5258
5259        // decode first page header
5260        let first_page = &file[4..];
5261        let mut prot = ThriftSliceInputProtocol::new(first_page);
5262        let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5263        let stats = hdr.data_page_header.unwrap().statistics;
5264        assert!(stats.is_some());
5265        let stats = stats.unwrap();
5266        // check that min/max were properly truncated
5267        assert!(!stats.is_max_value_exact.unwrap());
5268        assert!(!stats.is_min_value_exact.unwrap());
5269        assert_eq!(stats.max_value.unwrap(), b"Bm");
5270        assert_eq!(stats.min_value.unwrap(), b"Bl");
5271
5272        // check second page now
5273        let second_page = &prot.as_slice()[hdr.compressed_page_size as usize..];
5274        let mut prot = ThriftSliceInputProtocol::new(second_page);
5275        let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5276        let stats = hdr.data_page_header.unwrap().statistics;
5277        assert!(stats.is_some());
5278        let stats = stats.unwrap();
5279        // check that min/max were properly truncated
5280        assert!(!stats.is_max_value_exact.unwrap());
5281        assert!(!stats.is_min_value_exact.unwrap());
5282        assert_eq!(stats.max_value.unwrap(), b"Bm");
5283        assert_eq!(stats.min_value.unwrap(), b"Bl");
5284    }
5285
5286    #[test]
5287    fn test_page_encoding_statistics_roundtrip() {
5288        let batch_schema = Schema::new(vec![Field::new(
5289            "int32",
5290            arrow_schema::DataType::Int32,
5291            false,
5292        )]);
5293
5294        let batch = RecordBatch::try_new(
5295            Arc::new(batch_schema.clone()),
5296            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5297        )
5298        .unwrap();
5299
5300        let mut file: File = tempfile::tempfile().unwrap();
5301        let mut writer = ArrowWriter::try_new(&mut file, Arc::new(batch_schema), None).unwrap();
5302        writer.write(&batch).unwrap();
5303        let file_metadata = writer.close().unwrap();
5304
5305        assert_eq!(file_metadata.num_row_groups(), 1);
5306        assert_eq!(file_metadata.row_group(0).num_columns(), 1);
5307        assert!(
5308            file_metadata
5309                .row_group(0)
5310                .column(0)
5311                .page_encoding_stats()
5312                .is_some()
5313        );
5314        let chunk_page_stats = file_metadata
5315            .row_group(0)
5316            .column(0)
5317            .page_encoding_stats()
5318            .unwrap();
5319
5320        // check that the read metadata is also correct
5321        let options = ReadOptionsBuilder::new()
5322            .with_page_index()
5323            .with_encoding_stats_as_mask(false)
5324            .build();
5325        let reader = SerializedFileReader::new_with_options(file, options).unwrap();
5326
5327        let rowgroup = reader.get_row_group(0).expect("row group missing");
5328        assert_eq!(rowgroup.num_columns(), 1);
5329        let column = rowgroup.metadata().column(0);
5330        assert!(column.page_encoding_stats().is_some());
5331        let file_page_stats = column.page_encoding_stats().unwrap();
5332        assert_eq!(chunk_page_stats, file_page_stats);
5333    }
5334
5335    #[test]
5336    #[cfg_attr(miri, ignore)] // Takes too long
5337    fn test_different_dict_page_size_limit() {
5338        let array = Arc::new(Int64Array::from_iter(0..1024 * 1024));
5339        let schema = Arc::new(Schema::new(vec![
5340            Field::new("col0", arrow_schema::DataType::Int64, false),
5341            Field::new("col1", arrow_schema::DataType::Int64, false),
5342        ]));
5343        let batch =
5344            arrow_array::RecordBatch::try_new(schema.clone(), vec![array.clone(), array]).unwrap();
5345
5346        let props = WriterProperties::builder()
5347            .set_dictionary_page_size_limit(1024 * 1024)
5348            .set_column_dictionary_page_size_limit(ColumnPath::from("col1"), 1024 * 1024 * 4)
5349            .build();
5350        let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5351        writer.write(&batch).unwrap();
5352        let data = Bytes::from(writer.into_inner().unwrap());
5353
5354        let mut metadata = ParquetMetaDataReader::new();
5355        metadata.try_parse(&data).unwrap();
5356        let metadata = metadata.finish().unwrap();
5357        let col0_meta = metadata.row_group(0).column(0);
5358        let col1_meta = metadata.row_group(0).column(1);
5359
5360        let get_dict_page_size = move |meta: &ColumnChunkMetaData| {
5361            let mut reader =
5362                SerializedPageReader::new(Arc::new(data.clone()), meta, 0, None).unwrap();
5363            let page = reader.get_next_page().unwrap().unwrap();
5364            match page {
5365                Page::DictionaryPage { buf, .. } => buf.len(),
5366                _ => panic!("expected DictionaryPage"),
5367            }
5368        };
5369
5370        assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
5371        assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
5372    }
5373
5374    #[test]
5375    #[cfg_attr(miri, ignore)] // Takes too long
5376    fn test_arrow_writer_granular_mode_roundtrip() {
5377        // Granular mode subdivides chunks and writes more pages than the
5378        // default batched path. Make sure the data we write back is
5379        // bit-identical to what went in — page-count assertions elsewhere
5380        // only prove pages were cut, not that the encoded data is correct.
5381        //
5382        // Mix value sizes so that the cumulative-byte-budget cutoff
5383        // lands mid-chunk, exercising both batched and granular paths
5384        // within the same `write_batch_internal` call.
5385        let small = "tiny".to_string();
5386        let big = "x".repeat(64 * 1024);
5387        let strings: Vec<String> = (0..256)
5388            .map(|i| {
5389                if i % 16 == 0 {
5390                    big.clone()
5391                } else {
5392                    small.clone()
5393                }
5394            })
5395            .collect();
5396
5397        let schema = Arc::new(Schema::new(vec![Field::new(
5398            "col",
5399            ArrowDataType::Utf8,
5400            false,
5401        )]));
5402        let batch = RecordBatch::try_new(
5403            schema.clone(),
5404            vec![Arc::new(StringArray::from(strings.clone())) as _],
5405        )
5406        .unwrap();
5407
5408        let props = WriterProperties::builder()
5409            .set_dictionary_enabled(false)
5410            .set_data_page_size_limit(16 * 1024)
5411            .build();
5412        let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5413        writer.write(&batch).unwrap();
5414        let data = Bytes::from(writer.into_inner().unwrap());
5415
5416        let mut reader = ParquetRecordBatchReader::try_new(data, 1024).unwrap();
5417        let read = reader.next().unwrap().unwrap();
5418        assert!(reader.next().is_none(), "expected one batch");
5419        let col = read
5420            .column(0)
5421            .as_any()
5422            .downcast_ref::<StringArray>()
5423            .unwrap();
5424        assert_eq!(col.len(), strings.len());
5425        for (i, expected) in strings.iter().enumerate() {
5426            assert_eq!(
5427                col.value(i),
5428                expected.as_str(),
5429                "value mismatch at index {i}"
5430            );
5431        }
5432    }
5433
5434    #[test]
5435    fn test_arrow_writer_all_null_string_column() {
5436        // The `LevelDataRef::value_count` Uniform branch with
5437        // `value != max_def` (entirely-null chunk) must return 0 so the
5438        // sub-batch sizer short-circuits to batch mode without trying
5439        // to estimate byte budgets for non-existent values.
5440        let num_rows = 1024;
5441        let schema = Arc::new(Schema::new(vec![Field::new(
5442            "col",
5443            ArrowDataType::Utf8,
5444            true,
5445        )]));
5446        let nulls: Vec<Option<&str>> = vec![None; num_rows];
5447        let batch = RecordBatch::try_new(
5448            schema.clone(),
5449            vec![Arc::new(StringArray::from(nulls)) as _],
5450        )
5451        .unwrap();
5452
5453        let props = WriterProperties::builder()
5454            .set_dictionary_enabled(false)
5455            .set_data_page_size_limit(16 * 1024)
5456            .build();
5457        let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5458        writer.write(&batch).unwrap();
5459        let data = Bytes::from(writer.into_inner().unwrap());
5460
5461        // Re-parse the file: row group has one column, every row is
5462        // null, all data pages report `num_rows / page_count` rows.
5463        let mut metadata = ParquetMetaDataReader::new();
5464        metadata.try_parse(&data).unwrap();
5465        let metadata = metadata.finish().unwrap();
5466        let row_group = metadata.row_group(0);
5467        let col_meta = row_group.column(0);
5468        assert_eq!(row_group.num_rows() as usize, num_rows);
5469        // Statistics record `null_count = num_rows` — proves every value
5470        // was written as null.
5471        if let Some(stats) = col_meta.statistics() {
5472            assert_eq!(
5473                stats.null_count_opt().unwrap_or(0) as usize,
5474                num_rows,
5475                "expected all-null column to report null_count = num_rows"
5476            );
5477        }
5478
5479        let mut reader =
5480            SerializedPageReader::new(Arc::new(data.clone()), col_meta, num_rows, None).unwrap();
5481        let mut total_values = 0u32;
5482        while let Some(page) = reader.get_next_page().unwrap() {
5483            if matches!(page, Page::DataPage { .. } | Page::DataPageV2 { .. }) {
5484                total_values += page.num_values();
5485            }
5486        }
5487        assert_eq!(
5488            total_values as usize, num_rows,
5489            "expected every level position to be represented in some page"
5490        );
5491    }
5492
5493    struct WriteBatchesShape {
5494        num_batches: usize,
5495        rows_per_batch: usize,
5496        row_size: usize,
5497    }
5498
5499    /// Helper function to write batches with the provided `WriteBatchesShape` into an `ArrowWriter`
5500    fn write_batches(
5501        WriteBatchesShape {
5502            num_batches,
5503            rows_per_batch,
5504            row_size,
5505        }: WriteBatchesShape,
5506        props: WriterProperties,
5507    ) -> ParquetRecordBatchReaderBuilder<File> {
5508        let schema = Arc::new(Schema::new(vec![Field::new(
5509            "str",
5510            ArrowDataType::Utf8,
5511            false,
5512        )]));
5513        let file = tempfile::tempfile().unwrap();
5514        let mut writer =
5515            ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
5516
5517        for batch_idx in 0..num_batches {
5518            let strings: Vec<String> = (0..rows_per_batch)
5519                .map(|i| format!("{:0>width$}", batch_idx * 10 + i, width = row_size))
5520                .collect();
5521            let array = StringArray::from(strings);
5522            let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
5523            writer.write(&batch).unwrap();
5524        }
5525        writer.close().unwrap();
5526        ParquetRecordBatchReaderBuilder::try_new(file).unwrap()
5527    }
5528
5529    #[test]
5530    // When both limits are None, all data should go into a single row group
5531    fn test_row_group_limit_none_writes_single_row_group() {
5532        let props = WriterProperties::builder()
5533            .set_max_row_group_row_count(None)
5534            .set_max_row_group_bytes(None)
5535            .build();
5536
5537        let builder = write_batches(
5538            WriteBatchesShape {
5539                num_batches: 1,
5540                rows_per_batch: 1000,
5541                row_size: 4,
5542            },
5543            props,
5544        );
5545
5546        assert_eq!(
5547            &row_group_sizes(builder.metadata()),
5548            &[1000],
5549            "With no limits, all rows should be in a single row group"
5550        );
5551    }
5552
5553    #[test]
5554    // When only max_row_group_size is set, respect the row limit
5555    fn test_row_group_limit_rows_only() {
5556        let props = WriterProperties::builder()
5557            .set_max_row_group_row_count(Some(300))
5558            .set_max_row_group_bytes(None)
5559            .build();
5560
5561        let builder = write_batches(
5562            WriteBatchesShape {
5563                num_batches: 1,
5564                rows_per_batch: 1000,
5565                row_size: 4,
5566            },
5567            props,
5568        );
5569
5570        assert_eq!(
5571            &row_group_sizes(builder.metadata()),
5572            &[300, 300, 300, 100],
5573            "Row groups should be split by row count"
5574        );
5575    }
5576
5577    #[test]
5578    #[cfg_attr(miri, ignore)] // Takes too long
5579    // A row limit far smaller than the batch splits it many times over; the split must not
5580    // consume stack proportional to the number of row groups.
5581    fn test_row_group_limit_rows_only_many_splits() {
5582        let props = WriterProperties::builder()
5583            .set_max_row_group_row_count(Some(1))
5584            .set_max_row_group_bytes(None)
5585            .build();
5586
5587        let rows = 50_000;
5588        let builder = write_batches(
5589            WriteBatchesShape {
5590                num_batches: 1,
5591                rows_per_batch: rows,
5592                row_size: 4,
5593            },
5594            props,
5595        );
5596
5597        let sizes = row_group_sizes(builder.metadata());
5598        assert_eq!(sizes.len(), rows, "Every row should get its own row group");
5599        assert_eq!(
5600            sizes.iter().sum::<i64>(),
5601            rows as i64,
5602            "Total rows should be preserved"
5603        );
5604    }
5605
5606    #[test]
5607    // When only max_row_group_bytes is set, respect the byte limit
5608    fn test_row_group_limit_bytes_only() {
5609        let props = WriterProperties::builder()
5610            .set_max_row_group_row_count(None)
5611            // Set byte limit to approximately fit ~30 rows worth of data (~100 bytes each)
5612            .set_max_row_group_bytes(Some(3500))
5613            .build();
5614
5615        let builder = write_batches(
5616            WriteBatchesShape {
5617                num_batches: 10,
5618                rows_per_batch: 10,
5619                row_size: 100,
5620            },
5621            props,
5622        );
5623
5624        let sizes = row_group_sizes(builder.metadata());
5625
5626        assert!(
5627            sizes.len() > 1,
5628            "Should have multiple row groups due to byte limit, got {sizes:?}",
5629        );
5630
5631        let total_rows: i64 = sizes.iter().sum();
5632        assert_eq!(total_rows, 100, "Total rows should be preserved");
5633    }
5634
5635    #[test]
5636    // If an in-progress row group is already oversized, it should be flushed before writing more.
5637    fn test_row_group_limit_bytes_flushes_when_current_group_already_too_large() {
5638        let schema = Arc::new(Schema::new(vec![Field::new(
5639            "str",
5640            ArrowDataType::Utf8,
5641            false,
5642        )]));
5643        let file = tempfile::tempfile().unwrap();
5644
5645        // Start with no byte limit so we can intentionally build an oversized in-progress row group.
5646        let props = WriterProperties::builder()
5647            .set_max_row_group_row_count(None)
5648            .set_max_row_group_bytes(None)
5649            .build();
5650        let mut writer =
5651            ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
5652
5653        let first_array = StringArray::from(
5654            (0..10)
5655                .map(|i| format!("{i:0>100}"))
5656                .collect::<Vec<String>>(),
5657        );
5658        let first_batch =
5659            RecordBatch::try_new(schema.clone(), vec![Arc::new(first_array)]).unwrap();
5660        writer.write(&first_batch).unwrap();
5661        assert_eq!(writer.in_progress_rows(), 10);
5662
5663        // Tighten the limit below the current in-progress bytes to exercise:
5664        // `if current_bytes >= max_bytes { self.flush()?; ... }`
5665        writer.max_row_group_bytes = Some(1);
5666
5667        let second_array = StringArray::from(vec!["x".to_string()]);
5668        let second_batch =
5669            RecordBatch::try_new(schema.clone(), vec![Arc::new(second_array)]).unwrap();
5670        writer.write(&second_batch).unwrap();
5671        writer.close().unwrap();
5672        let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
5673
5674        assert_eq!(
5675            &row_group_sizes(builder.metadata()),
5676            &[10, 1],
5677            "The second write should flush an oversized in-progress row group first",
5678        );
5679    }
5680
5681    #[test]
5682    // When both limits are set, the row limit triggers first
5683    fn test_row_group_limit_both_row_wins_single_batch() {
5684        let props = WriterProperties::builder()
5685            .set_max_row_group_row_count(Some(200)) // Will trigger at 200 rows
5686            .set_max_row_group_bytes(Some(1024 * 1024)) // 1MB - won't trigger for small int data
5687            .build();
5688
5689        let builder = write_batches(
5690            WriteBatchesShape {
5691                num_batches: 1,
5692                row_size: 4,
5693                rows_per_batch: 1000,
5694            },
5695            props,
5696        );
5697
5698        assert_eq!(
5699            &row_group_sizes(builder.metadata()),
5700            &[200, 200, 200, 200, 200],
5701            "Row limit should trigger before byte limit"
5702        );
5703    }
5704
5705    #[test]
5706    // When both limits are set, the row limit triggers first
5707    fn test_row_group_limit_both_row_wins_multiple_batches() {
5708        let props = WriterProperties::builder()
5709            .set_max_row_group_row_count(Some(5)) // Will trigger every 5 rows
5710            .set_max_row_group_bytes(Some(9999)) // Won't trigger
5711            .build();
5712
5713        let builder = write_batches(
5714            WriteBatchesShape {
5715                num_batches: 10,
5716                rows_per_batch: 10,
5717                row_size: 100,
5718            },
5719            props,
5720        );
5721
5722        assert_eq!(
5723            &row_group_sizes(builder.metadata()),
5724            &[5; 20],
5725            "Row limit should trigger before byte limit"
5726        );
5727    }
5728
5729    #[test]
5730    // When both limits are set, the byte limit triggers first
5731    fn test_row_group_limit_both_bytes_wins() {
5732        let props = WriterProperties::builder()
5733            .set_max_row_group_row_count(Some(1000)) // Won't trigger for 100 rows
5734            .set_max_row_group_bytes(Some(3500)) // Will trigger at ~30-35 rows
5735            .build();
5736
5737        let builder = write_batches(
5738            WriteBatchesShape {
5739                num_batches: 10,
5740                rows_per_batch: 10,
5741                row_size: 100,
5742            },
5743            props,
5744        );
5745
5746        let sizes = row_group_sizes(builder.metadata());
5747
5748        assert!(
5749            sizes.len() > 1,
5750            "Byte limit should trigger before row limit, got {sizes:?}",
5751        );
5752
5753        assert!(
5754            sizes.iter().all(|&s| s < 1000),
5755            "No row group should hit the row limit"
5756        );
5757
5758        let total_rows: i64 = sizes.iter().sum();
5759        assert_eq!(total_rows, 100, "Total rows should be preserved");
5760    }
5761
5762    #[test]
5763    // Both limits can apply to the same batch: the row limit trims it to 5 rows, and the
5764    // byte limit then trims those 5 down to 4.
5765    fn test_row_group_limit_both_apply_to_same_batch() {
5766        let props = WriterProperties::builder()
5767            .set_max_row_group_row_count(Some(15))
5768            .set_max_row_group_bytes(Some(1500))
5769            .build();
5770
5771        let builder = write_batches(
5772            WriteBatchesShape {
5773                num_batches: 2,
5774                rows_per_batch: 10,
5775                row_size: 100,
5776            },
5777            props,
5778        );
5779
5780        assert_eq!(
5781            &row_group_sizes(builder.metadata()),
5782            &[14, 6],
5783            "Byte limit should still apply to a batch the row limit already split"
5784        );
5785    }
5786
5787    #[test]
5788    // A batch far larger than the byte limit must be split even when it is the first one
5789    // written to a row group, where there are no buffered rows to estimate the row size from.
5790    fn test_row_group_limit_bytes_splits_first_batch() {
5791        let props = WriterProperties::builder()
5792            .set_max_row_group_row_count(None)
5793            .set_max_row_group_bytes(Some(3500))
5794            .build();
5795
5796        let builder = write_batches(
5797            WriteBatchesShape {
5798                num_batches: 3,
5799                rows_per_batch: 100,
5800                row_size: 100,
5801            },
5802            props,
5803        );
5804
5805        let sizes = row_group_sizes(builder.metadata());
5806        assert!(
5807            sizes.iter().all(|&s| s < 100),
5808            "No row group should hold a whole batch, got {sizes:?}",
5809        );
5810        assert_eq!(
5811            sizes.iter().sum::<i64>(),
5812            300,
5813            "Total rows should be preserved"
5814        );
5815        for row_group in builder.metadata().row_groups() {
5816            assert!(
5817                row_group.total_byte_size() <= 3500 + 100,
5818                "Row groups should stay within one row of the byte limit, got {}",
5819                row_group.total_byte_size(),
5820            );
5821        }
5822    }
5823
5824    #[test]
5825    fn arrow_column_chunk_close_mut_drops_column_index() {
5826        use crate::arrow::ArrowSchemaConverter;
5827        use crate::file::writer::SerializedFileWriter;
5828
5829        let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)]));
5830        let props = Arc::new(
5831            WriterProperties::builder()
5832                .set_statistics_enabled(EnabledStatistics::Page)
5833                .build(),
5834        );
5835        let parquet_schema = ArrowSchemaConverter::new()
5836            .with_coerce_types(props.coerce_types())
5837            .convert(&schema)
5838            .unwrap();
5839
5840        let mut buf = Vec::with_capacity(1024);
5841        let mut writer =
5842            SerializedFileWriter::new(&mut buf, parquet_schema.root_schema_ptr(), props.clone())
5843                .unwrap();
5844
5845        let factory = ArrowRowGroupWriterFactory::try_new(&writer, Arc::clone(&schema)).unwrap();
5846        let mut col_writers = factory.create_column_writers(0).unwrap();
5847        let arr: ArrayRef = Arc::new(Int32Array::from_iter_values(0..64));
5848        for leaves in compute_leaves(schema.field(0), &arr).unwrap() {
5849            col_writers[0].write(&leaves).unwrap();
5850        }
5851        let mut chunk = col_writers.pop().unwrap().close().unwrap();
5852
5853        // Immutable accessor exposes the close result produced at close time.
5854        assert!(
5855            chunk.close().column_index.is_some(),
5856            "EnabledStatistics::Page should produce a column_index"
5857        );
5858
5859        // Mutable accessor lets callers drop the page-level index before append.
5860        chunk.close_mut().column_index = None;
5861        assert!(chunk.close().column_index.is_none());
5862
5863        let mut rg = writer.next_row_group().unwrap();
5864        chunk.append_to_row_group(&mut rg).unwrap();
5865        rg.close().unwrap();
5866        let file_meta = writer.close().unwrap();
5867
5868        // After dropping column_index, the resulting file records no column
5869        // index offset/length for this chunk.
5870        let cc = file_meta.row_group(0).column(0);
5871        assert!(cc.column_index_range().is_none());
5872    }
5873
5874    #[test]
5875    #[cfg_attr(miri, ignore)] // Takes too long
5876    fn test_number_distinct_values_exact_count() {
5877        // 50 distinct Int32 values repeated across 100k rows, with every 7th row null.
5878        // Nulls must not be counted as a distinct value.
5879        let cardinality = 50u32;
5880        let array: ArrayRef = Arc::new(Int32Array::from_iter((0..100_000u32).map(|i| {
5881            if i % 7 == 0 {
5882                None
5883            } else {
5884                Some((i % cardinality) as i32)
5885            }
5886        })));
5887        let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, true)]));
5888        let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
5889
5890        let props = WriterProperties::builder()
5891            .set_write_row_group_number_distinct_values(true)
5892            .build();
5893        let mut buf = Vec::new();
5894        let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), Some(props)).unwrap();
5895        writer.write(&batch).unwrap();
5896        let metadata = writer.close().unwrap();
5897
5898        let count = metadata
5899            .row_group(0)
5900            .column(0)
5901            .statistics()
5902            .and_then(|s| s.distinct_count_opt())
5903            .expect("distinct_count should be set");
5904        // Must equal cardinality exactly; nulls must not inflate the count.
5905        assert_eq!(count, cardinality as u64);
5906    }
5907
5908    #[test]
5909    fn test_number_distinct_values_view_types() {
5910        // 5 distinct values repeated across 30 rows, with every 4th row null.
5911        // Verifies Utf8View is counted correctly (BinaryView shares the same code path).
5912        let cardinality = 5u32;
5913        let distinct_strings = ["alpha", "beta", "gamma", "delta", "epsilon"];
5914
5915        let string_view_col: ArrayRef = Arc::new(StringViewArray::from_iter((0..30u32).map(|i| {
5916            if i % 4 == 0 {
5917                None
5918            } else {
5919                Some(distinct_strings[(i % cardinality) as usize])
5920            }
5921        })));
5922
5923        let schema = Arc::new(Schema::new(vec![Field::new(
5924            "string_view_col",
5925            DataType::Utf8View,
5926            true,
5927        )]));
5928        let batch = RecordBatch::try_new(schema, vec![string_view_col]).unwrap();
5929
5930        let props = WriterProperties::builder()
5931            .set_write_row_group_number_distinct_values(true)
5932            .build();
5933        let mut parquet_bytes = Vec::new();
5934        let mut writer =
5935            ArrowWriter::try_new(&mut parquet_bytes, batch.schema(), Some(props)).unwrap();
5936        writer.write(&batch).unwrap();
5937        let metadata = writer.close().unwrap();
5938
5939        let distinct_count = metadata
5940            .row_group(0)
5941            .column(0)
5942            .statistics()
5943            .and_then(|s| s.distinct_count_opt())
5944            .expect("distinct_count should be set for Utf8View column");
5945        assert_eq!(distinct_count, cardinality as u64);
5946    }
5947
5948    #[test]
5949    fn test_number_distinct_values_not_written_by_default() {
5950        let array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..100));
5951        let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, false)]));
5952        let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
5953
5954        let mut buf = Vec::new();
5955        let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), None).unwrap();
5956        writer.write(&batch).unwrap();
5957        let metadata = writer.close().unwrap();
5958
5959        let count = metadata
5960            .row_group(0)
5961            .column(0)
5962            .statistics()
5963            .and_then(|s| s.distinct_count_opt());
5964        assert!(count.is_none());
5965    }
5966
5967    #[test]
5968    fn test_dictionary_ndv_single_batch() {
5969        // Dictionary array with 3 distinct string values repeated many times.
5970        // NDV must equal the number of distinct values in the dictionary (3),
5971        // not the number of rows.
5972        let keys = Int32Array::from(vec![0, 1, 2, 0, 1, 2, 0, 1, 2]);
5973        let values: ArrayRef = Arc::new(StringArray::from(vec!["cat", "dog", "bird"]));
5974        let dict: ArrayRef = Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap());
5975
5976        let schema = Arc::new(Schema::new(vec![Field::new(
5977            "x",
5978            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
5979            false,
5980        )]));
5981        let batch = RecordBatch::try_new(schema, vec![dict]).unwrap();
5982
5983        let props = WriterProperties::builder()
5984            .set_write_row_group_number_distinct_values(true)
5985            .build();
5986        let mut buf = Vec::new();
5987        let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), Some(props)).unwrap();
5988        writer.write(&batch).unwrap();
5989        let metadata = writer.close().unwrap();
5990
5991        let count = metadata
5992            .row_group(0)
5993            .column(0)
5994            .statistics()
5995            .and_then(|s| s.distinct_count_opt())
5996            .expect("distinct_count should be set");
5997        assert_eq!(count, 3);
5998    }
5999
6000    #[test]
6001    fn test_dictionary_ndv_excludes_unreferenced_values() {
6002        // Keys only reference indices 0 and 1; value at index 2 ("unreferenced") should not
6003        // count toward NDV even though it appears in the dictionary's values array.
6004        let keys = Int32Array::from(vec![0, 1, 0, 1]);
6005        let values: ArrayRef = Arc::new(StringArray::from(vec!["cat", "dog", "unreferenced"]));
6006        let dict: ArrayRef = Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap());
6007
6008        let schema = Arc::new(Schema::new(vec![Field::new(
6009            "x",
6010            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
6011            false,
6012        )]));
6013        let batch = RecordBatch::try_new(schema, vec![dict]).unwrap();
6014
6015        let props = WriterProperties::builder()
6016            .set_write_row_group_number_distinct_values(true)
6017            .build();
6018        let mut buf = Vec::new();
6019        let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), Some(props)).unwrap();
6020        writer.write(&batch).unwrap();
6021        let metadata = writer.close().unwrap();
6022
6023        let count = metadata
6024            .row_group(0)
6025            .column(0)
6026            .statistics()
6027            .and_then(|s| s.distinct_count_opt())
6028            .expect("distinct_count should be set");
6029        assert_eq!(
6030            count, 2,
6031            "unreferenced dictionary values must not count toward NDV"
6032        );
6033    }
6034
6035    #[test]
6036    fn test_dictionary_ndv_across_batches_regression() {
6037        // Regression test for https://github.com/apache/arrow-rs/issues/11172.
6038        let make_dict_batch = |a: &str, b: &str| -> RecordBatch {
6039            let keys = Int32Array::from(vec![0, 1, 0, 1]);
6040            let values: ArrayRef = Arc::new(StringArray::from(vec![a, b]));
6041            let dict: ArrayRef =
6042                Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap());
6043            let schema = Arc::new(Schema::new(vec![Field::new(
6044                "x",
6045                DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
6046                false,
6047            )]));
6048            RecordBatch::try_new(schema, vec![dict]).unwrap()
6049        };
6050
6051        // batch1: dict = ["cat", "dog"], batch2: dict = ["fish", "cat"]
6052        // Distinct values across both batches: "cat", "dog", "fish" NDV = 3
6053        let batch1 = make_dict_batch("cat", "dog");
6054        let batch2 = make_dict_batch("fish", "cat");
6055
6056        let props = WriterProperties::builder()
6057            .set_write_row_group_number_distinct_values(true)
6058            .build();
6059        let mut buf = Vec::new();
6060        let mut writer = ArrowWriter::try_new(&mut buf, batch1.schema(), Some(props)).unwrap();
6061        writer.write(&batch1).unwrap();
6062        writer.write(&batch2).unwrap();
6063        let metadata = writer.close().unwrap();
6064
6065        let count = metadata
6066            .row_group(0)
6067            .column(0)
6068            .statistics()
6069            .and_then(|s| s.distinct_count_opt())
6070            .expect("distinct_count should be set");
6071        assert_eq!(
6072            count, 3,
6073            "NDV should count distinct values, not distinct key indices"
6074        );
6075    }
6076
6077    #[test]
6078    fn ree_struct_with_ree_child() {
6079        // Struct with a REE string field and a REE int field — confirms
6080        // recursion visits every child and each collapses to the right leaf type.
6081        let run_ends = Int32Array::from(vec![2i32, 3, 5]);
6082
6083        let col_a: ArrayRef = Arc::new(
6084            RunArray::try_new(
6085                &run_ends,
6086                &StringArray::from(vec![Some("foo"), None, Some("bar")]),
6087            )
6088            .unwrap(),
6089        );
6090        let col_b: ArrayRef = Arc::new(
6091            RunArray::try_new(&run_ends, &Int32Array::from(vec![Some(1), None, Some(2)])).unwrap(),
6092        );
6093
6094        let struct_array: ArrayRef = Arc::new(StructArray::new(
6095            Fields::from(vec![
6096                Field::new("a", col_a.data_type().clone(), true),
6097                Field::new("b", col_b.data_type().clone(), true),
6098            ]),
6099            vec![col_a, col_b],
6100            None,
6101        ));
6102
6103        let schema = Arc::new(Schema::new(vec![Field::new(
6104            "row",
6105            struct_array.data_type().clone(),
6106            true,
6107        )]));
6108        let batch = RecordBatch::try_new(schema.clone(), vec![struct_array]).unwrap();
6109
6110        let mut buf = Vec::new();
6111        let mut writer = ArrowWriter::try_new(&mut buf, schema, None).unwrap();
6112        writer.write(&batch).unwrap();
6113        let metadata = writer.close().unwrap();
6114
6115        let parquet_schema = metadata.file_metadata().schema_descr();
6116        assert_eq!(parquet_schema.num_columns(), 2);
6117        assert_eq!(
6118            parquet_schema.column(0).physical_type(),
6119            crate::basic::Type::BYTE_ARRAY
6120        );
6121        assert_eq!(parquet_schema.column(0).path().string(), "row.a");
6122        assert_eq!(
6123            parquet_schema.column(1).physical_type(),
6124            crate::basic::Type::INT32
6125        );
6126        assert_eq!(parquet_schema.column(1).path().string(), "row.b");
6127    }
6128
6129    /// Regression test for https://github.com/apache/arrow-rs/issues/11073
6130    ///
6131    /// This checks that the `ArrowColumnWriter` path produces the same min/max
6132    /// statistics as writing the column directly with the low-level
6133    /// `SerializedFileWriter` API, for both `BinaryArray` (the buggy case)
6134    /// and `FixedSizeBinaryArray` (which already worked correctly).
6135    #[test]
6136    fn test_decimal_byte_array_statistics_signed_comparison() {
6137        use crate::basic::{LogicalType, Repetition, Type as PhysicalType};
6138        use crate::data_type::{ByteArray, ByteArrayType};
6139        use crate::schema::types::Type as SchemaType;
6140
6141        // 1-byte two's-complement decimal values: -1, 0, 1
6142        const VALUES: [&[u8]; 3] = [&[0xffu8], &[0x00], &[0x01]];
6143
6144        fn new_file(physical: PhysicalType) -> SerializedFileWriter<Vec<u8>> {
6145            let mut field = SchemaType::primitive_type_builder("value", physical)
6146                .with_repetition(Repetition::REQUIRED)
6147                .with_logical_type(Some(LogicalType::decimal(0, 2)))
6148                .with_precision(2)
6149                .with_scale(0);
6150            if physical == PhysicalType::FIXED_LEN_BYTE_ARRAY {
6151                field = field.with_length(1);
6152            }
6153            let root = SchemaType::group_type_builder("root")
6154                .with_fields(vec![Arc::new(field.build().unwrap())])
6155                .build()
6156                .unwrap();
6157            SerializedFileWriter::new(Vec::new(), Arc::new(root), Default::default()).unwrap()
6158        }
6159
6160        fn read_minmax(file: Vec<u8>) -> (i8, i8) {
6161            let reader = SerializedFileReader::new(Bytes::from(file)).unwrap();
6162            let metadata = reader.metadata();
6163            assert_eq!(metadata.file_metadata().num_rows(), 3);
6164            let stats = metadata
6165                .row_group(0)
6166                .column(0)
6167                .statistics()
6168                .expect("missing statistics");
6169            (
6170                stats.min_bytes_opt().unwrap()[0] as i8,
6171                stats.max_bytes_opt().unwrap()[0] as i8,
6172            )
6173        }
6174
6175        fn direct_parquet() -> Vec<u8> {
6176            let mut file = new_file(PhysicalType::BYTE_ARRAY);
6177            let mut group = file.next_row_group().unwrap();
6178            let mut column = group.next_column().unwrap().unwrap();
6179            let values = VALUES
6180                .iter()
6181                .map(|v| ByteArray::from(v.to_vec()))
6182                .collect::<Vec<_>>();
6183            column
6184                .typed::<ByteArrayType>()
6185                .write_batch(&values, None, None)
6186                .unwrap();
6187            column.close().unwrap();
6188            group.close().unwrap();
6189            file.into_inner().unwrap()
6190        }
6191
6192        fn arrow_bounds(mut file: SerializedFileWriter<Vec<u8>>, array: ArrayRef) -> Vec<u8> {
6193            let field = Field::new("value", array.data_type().clone(), false);
6194            let schema = Arc::new(Schema::new(vec![field.clone()]));
6195            let factory = ArrowRowGroupWriterFactory::try_new(&file, schema).unwrap();
6196            let mut writer: ArrowColumnWriter = factory.create_column_writers(0).unwrap().remove(0);
6197            let leaf = compute_leaves(&field, &array).unwrap().remove(0);
6198            writer.write(&leaf).unwrap();
6199            let mut group = file.next_row_group().unwrap();
6200            writer
6201                .close()
6202                .unwrap()
6203                .append_to_row_group(&mut group)
6204                .unwrap();
6205            group.close().unwrap();
6206            file.into_inner().unwrap()
6207        }
6208
6209        fn arrow_binary() -> Vec<u8> {
6210            let file = new_file(PhysicalType::BYTE_ARRAY);
6211            let array: ArrayRef = Arc::new(BinaryArray::from_vec(VALUES.to_vec()));
6212            arrow_bounds(file, array)
6213        }
6214
6215        fn arrow_fixed_size_binary() -> Vec<u8> {
6216            let file = new_file(PhysicalType::FIXED_LEN_BYTE_ARRAY);
6217            let array: ArrayRef =
6218                Arc::new(FixedSizeBinaryArray::try_from_iter(VALUES.into_iter()).unwrap());
6219            arrow_bounds(file, array)
6220        }
6221
6222        let direct = read_minmax(direct_parquet());
6223        let fixed = read_minmax(arrow_fixed_size_binary());
6224        let binary = read_minmax(arrow_binary());
6225
6226        assert_eq!(direct, (-1, 1), "incorrect direct Parquet stats");
6227        assert_eq!(fixed, (-1, 1), "incorrect Arrow FixedSizeBinaryArray stats");
6228        assert_eq!(binary, (-1, 1), "incorrect Arrow BinaryArray stats");
6229    }
6230}