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, and lets
1803/// [`ArrowSbbf`](crate::arrow::bloom_filter::ArrowSbbf) check converted values
1804/// one at a time without collecting them.
1805pub(crate) trait PrimitiveVisitor<T> {
1806    type Output;
1807
1808    /// Receives values that are stored as they are
1809    fn slice(self, values: &[T]) -> Result<Self::Output>;
1810
1811    /// Receives values that are each stored as `convert(value)`
1812    fn map<S: Copy>(self, values: &[S], convert: impl Fn(S) -> T) -> Result<Self::Output>;
1813}
1814
1815/// Passes the values [`ArrowWriter`] stores for `column` in an `INT32` column to `visitor`
1816pub(crate) fn visit_i32<V: PrimitiveVisitor<i32>>(
1817    column: &dyn arrow_array::Array,
1818    visitor: V,
1819) -> Result<V::Output> {
1820    // Note: this should match the contents of arrow_to_parquet_type
1821    match column.data_type() {
1822        ArrowDataType::Null => visitor.slice(Int32Array::new_null(column.len()).values()),
1823        ArrowDataType::Int8 => {
1824            visitor.map(column.as_primitive::<Int8Type>().values(), |x| x as i32)
1825        }
1826        ArrowDataType::Int16 => {
1827            visitor.map(column.as_primitive::<Int16Type>().values(), |x| x as i32)
1828        }
1829        ArrowDataType::Int32 => visitor.slice(column.as_primitive::<Int32Type>().values()),
1830        ArrowDataType::UInt8 => {
1831            visitor.map(column.as_primitive::<UInt8Type>().values(), |x| x as i32)
1832        }
1833        ArrowDataType::UInt16 => {
1834            visitor.map(column.as_primitive::<UInt16Type>().values(), |x| x as i32)
1835        }
1836        ArrowDataType::UInt32 => {
1837            // follow C++ implementation and use overflow/reinterpret cast from  u32 to i32 which will map
1838            // `(i32::MAX as u32)..u32::MAX` to `i32::MIN..0`
1839            let values = column.as_primitive::<UInt32Type>().values();
1840            visitor.slice(values.inner().typed_data())
1841        }
1842        ArrowDataType::Date32 => visitor.slice(column.as_primitive::<Date32Type>().values()),
1843        ArrowDataType::Time32(TimeUnit::Second) => {
1844            visitor.slice(column.as_primitive::<Time32SecondType>().values())
1845        }
1846        ArrowDataType::Time32(TimeUnit::Millisecond) => {
1847            visitor.slice(column.as_primitive::<Time32MillisecondType>().values())
1848        }
1849        ArrowDataType::Date64 => {
1850            // If the column is a Date64, we truncate it
1851            visitor.map(column.as_primitive::<Date64Type>().values(), |x| {
1852                (x / 86_400_000) as i32
1853            })
1854        }
1855        ArrowDataType::Decimal32(_, _) => {
1856            visitor.slice(column.as_primitive::<Decimal32Type>().values())
1857        }
1858        ArrowDataType::Decimal64(_, _) => {
1859            // use the int32 to represent the decimal with low precision
1860            visitor.map(column.as_primitive::<Decimal64Type>().values(), |v| {
1861                v as i32
1862            })
1863        }
1864        ArrowDataType::Decimal128(_, _) => {
1865            // use the int32 to represent the decimal with low precision
1866            visitor.map(column.as_primitive::<Decimal128Type>().values(), |v| {
1867                v as i32
1868            })
1869        }
1870        ArrowDataType::Decimal256(_, _) => {
1871            // use the int32 to represent the decimal with low precision
1872            visitor.map(column.as_primitive::<Decimal256Type>().values(), |v| {
1873                v.as_i128() as i32
1874            })
1875        }
1876        d => Err(ParquetError::General(format!("Cannot coerce {d} to I32"))),
1877    }
1878}
1879
1880/// Passes the values [`ArrowWriter`] stores for `column` in an `INT64` column to `visitor`
1881pub(crate) fn visit_i64<V: PrimitiveVisitor<i64>>(
1882    column: &dyn arrow_array::Array,
1883    visitor: V,
1884) -> Result<V::Output> {
1885    // Note: this should match the contents of arrow_to_parquet_type
1886    match column.data_type() {
1887        ArrowDataType::Date64 => visitor.slice(column.as_primitive::<Date64Type>().values()),
1888        ArrowDataType::Int64 => visitor.slice(column.as_primitive::<Int64Type>().values()),
1889        ArrowDataType::UInt64 => {
1890            // follow C++ implementation and use overflow/reinterpret cast from  u64 to i64 which will map
1891            // `(i64::MAX as u64)..u64::MAX` to `i64::MIN..0`
1892            let values = column.as_primitive::<UInt64Type>().values();
1893            visitor.slice(values.inner().typed_data())
1894        }
1895        ArrowDataType::Time64(TimeUnit::Microsecond) => {
1896            visitor.slice(column.as_primitive::<Time64MicrosecondType>().values())
1897        }
1898        ArrowDataType::Time64(TimeUnit::Nanosecond) => {
1899            visitor.slice(column.as_primitive::<Time64NanosecondType>().values())
1900        }
1901        ArrowDataType::Timestamp(unit, _) => match unit {
1902            TimeUnit::Second => {
1903                visitor.slice(column.as_primitive::<TimestampSecondType>().values())
1904            }
1905            TimeUnit::Millisecond => {
1906                visitor.slice(column.as_primitive::<TimestampMillisecondType>().values())
1907            }
1908            TimeUnit::Microsecond => {
1909                visitor.slice(column.as_primitive::<TimestampMicrosecondType>().values())
1910            }
1911            TimeUnit::Nanosecond => {
1912                visitor.slice(column.as_primitive::<TimestampNanosecondType>().values())
1913            }
1914        },
1915        ArrowDataType::Duration(unit) => match unit {
1916            TimeUnit::Second => visitor.slice(column.as_primitive::<DurationSecondType>().values()),
1917            TimeUnit::Millisecond => {
1918                visitor.slice(column.as_primitive::<DurationMillisecondType>().values())
1919            }
1920            TimeUnit::Microsecond => {
1921                visitor.slice(column.as_primitive::<DurationMicrosecondType>().values())
1922            }
1923            TimeUnit::Nanosecond => {
1924                visitor.slice(column.as_primitive::<DurationNanosecondType>().values())
1925            }
1926        },
1927        ArrowDataType::Decimal32(_, _) => {
1928            // ArrowWriter stores Decimal32 as INT32, but ArrowSbbf can check
1929            // Decimal32 values against an INT64 decimal column
1930            visitor.map(column.as_primitive::<Decimal32Type>().values(), |v| {
1931                v as i64
1932            })
1933        }
1934        ArrowDataType::Decimal64(_, _) => {
1935            visitor.slice(column.as_primitive::<Decimal64Type>().values())
1936        }
1937        ArrowDataType::Decimal128(_, _) => {
1938            // use the int64 to represent the decimal with low precision
1939            visitor.map(column.as_primitive::<Decimal128Type>().values(), |v| {
1940                v as i64
1941            })
1942        }
1943        ArrowDataType::Decimal256(_, _) => {
1944            // use the int64 to represent the decimal with low precision
1945            visitor.map(column.as_primitive::<Decimal256Type>().values(), |v| {
1946                v.as_i128() as i64
1947            })
1948        }
1949        d => Err(ParquetError::General(format!("Cannot coerce {d} to I64"))),
1950    }
1951}
1952
1953/// Writes the values passed to it to an `INT32` or `INT64` column
1954struct WritePrimitive<'a, 'b, E: ColumnValueEncoder> {
1955    writer: &'a mut GenericColumnWriter<'b, E>,
1956    levels: &'a ArrayLevels,
1957}
1958
1959impl<T, E> PrimitiveVisitor<T> for WritePrimitive<'_, '_, E>
1960where
1961    E: ColumnValueEncoder<Values = [T]>,
1962{
1963    type Output = usize;
1964
1965    fn slice(self, values: &[T]) -> Result<usize> {
1966        write_primitive(self.writer, values, self.levels)
1967    }
1968
1969    fn map<S: Copy>(self, values: &[S], convert: impl Fn(S) -> T) -> Result<usize> {
1970        let values: Vec<T> = values.iter().map(|v| convert(*v)).collect();
1971        write_primitive(self.writer, &values, self.levels)
1972    }
1973}
1974
1975/// Receives the values [`ArrowWriter`] stores for an Arrow array in a
1976/// `FIXED_LEN_BYTE_ARRAY` column, which are also the values hashed into the
1977/// column's bloom filter
1978///
1979/// [`visit_fixed_len_byte_array`] holds the conversion from each Arrow type
1980pub(crate) trait FixedLenByteArrayVisitor {
1981    type Output;
1982
1983    /// Returns the length of the values that decimals of `precision` are stored as
1984    ///
1985    /// `native_length` is the length of the Arrow decimal type, for example 16
1986    /// for `Decimal128`.
1987    fn decimal_length(&self, precision: u8, native_length: usize) -> usize;
1988
1989    /// Receives values that are `length` bytes long
1990    ///
1991    /// `fill(i, value)` writes the value at index `i` of the array to `value`,
1992    /// which is `length` zeroed bytes.
1993    fn values(self, length: usize, fill: impl Fn(usize, &mut [u8])) -> Result<Self::Output>;
1994}
1995
1996/// Passes the values [`ArrowWriter`] stores for `column` in a `FIXED_LEN_BYTE_ARRAY`
1997/// column to `visitor`
1998pub(crate) fn visit_fixed_len_byte_array<V: FixedLenByteArrayVisitor>(
1999    column: &dyn arrow_array::Array,
2000    visitor: V,
2001) -> Result<V::Output> {
2002    match column.data_type() {
2003        ArrowDataType::Interval(interval_unit) => match interval_unit {
2004            IntervalUnit::YearMonth => {
2005                // 12 bytes of months, days and milliseconds (4 bytes each). An Arrow
2006                // YearMonth interval only stores months, so only the first 4 bytes are set.
2007                let array = column.as_primitive::<IntervalYearMonthType>();
2008                visitor.values(12, |i, value| {
2009                    value[0..4].copy_from_slice(&array.value(i).to_le_bytes());
2010                })
2011            }
2012            IntervalUnit::DayTime => {
2013                // 12 bytes of months, days and milliseconds (4 bytes each). An Arrow
2014                // DayTime interval only stores days and milliseconds, so the first 4
2015                // bytes are not set.
2016                let array = column.as_primitive::<IntervalDayTimeType>();
2017                visitor.values(12, |i, value| {
2018                    let interval = array.value(i);
2019                    value[4..8].copy_from_slice(&interval.days.to_le_bytes());
2020                    value[8..12].copy_from_slice(&interval.milliseconds.to_le_bytes());
2021                })
2022            }
2023            IntervalUnit::MonthDayNano => Err(ParquetError::NYI(format!(
2024                "Attempting to write an Arrow interval type {interval_unit:?} to parquet that is not yet implemented"
2025            ))),
2026        },
2027        ArrowDataType::FixedSizeBinary(_) => {
2028            let array = column.as_fixed_size_binary();
2029            visitor.values(array.value_size(), |i, value| {
2030                value.copy_from_slice(array.value(i));
2031            })
2032        }
2033        ArrowDataType::Decimal32(_, _) => {
2034            visit_decimal(column.as_primitive::<Decimal32Type>(), visitor)
2035        }
2036        ArrowDataType::Decimal64(_, _) => {
2037            visit_decimal(column.as_primitive::<Decimal64Type>(), visitor)
2038        }
2039        ArrowDataType::Decimal128(_, _) => {
2040            visit_decimal(column.as_primitive::<Decimal128Type>(), visitor)
2041        }
2042        ArrowDataType::Decimal256(_, _) => {
2043            visit_decimal(column.as_primitive::<Decimal256Type>(), visitor)
2044        }
2045        ArrowDataType::Float16 => {
2046            let array = column.as_primitive::<Float16Type>();
2047            visitor.values(2, |i, value| {
2048                value.copy_from_slice(&array.value(i).to_le_bytes());
2049            })
2050        }
2051        _ => Err(ParquetError::NYI(
2052            "Attempting to write an Arrow type that is not yet implemented".to_string(),
2053        )),
2054    }
2055}
2056
2057/// Passes the big-endian two's complement bytes of the values of `array` to `visitor`
2058///
2059/// The values are resized to [`FixedLenByteArrayVisitor::decimal_length`] bytes:
2060/// sign-extended if that is longer than the Arrow decimal type, or truncated to
2061/// the low bytes if it is shorter. A value too large for the length is truncated,
2062/// as [`ArrowWriter`] has always done, since Arrow does not validate precision.
2063fn visit_decimal<T: NativeDecimalType, V: FixedLenByteArrayVisitor>(
2064    array: &PrimitiveArray<T>,
2065    visitor: V,
2066) -> Result<V::Output> {
2067    let length = visitor.decimal_length(array.precision(), T::BYTE_LENGTH);
2068    let be_bytes = |i| T::to_be_bytes(array.value(i));
2069    // Compare the lengths once, not for each value
2070    match length.checked_sub(T::BYTE_LENGTH) {
2071        // Same length, which lets the compiler inline the copy
2072        Some(0) => visitor.values(length, |i, value| {
2073            value.copy_from_slice(be_bytes(i).as_ref());
2074        }),
2075        Some(extension) => visitor.values(length, |i, value| {
2076            let bytes = be_bytes(i);
2077            value[..extension].fill(sign_byte(bytes.as_ref()));
2078            value[extension..].copy_from_slice(bytes.as_ref());
2079        }),
2080        None => visitor.values(length, |i, value| {
2081            value.copy_from_slice(&be_bytes(i).as_ref()[T::BYTE_LENGTH - length..]);
2082        }),
2083    }
2084}
2085
2086/// Returns the byte that sign-extends the big-endian two's complement integer `bytes`
2087fn sign_byte(bytes: &[u8]) -> u8 {
2088    if bytes.first().is_some_and(|b| (*b as i8) < 0) {
2089        0xFF
2090    } else {
2091        0
2092    }
2093}
2094
2095/// Writes the values passed to it to a `FIXED_LEN_BYTE_ARRAY` column
2096struct WriteFixedLenByteArray<'a, 'b> {
2097    writer: &'a mut ColumnWriterImpl<'b, FixedLenByteArrayType>,
2098    levels: &'a ArrayLevels,
2099}
2100
2101impl FixedLenByteArrayVisitor for WriteFixedLenByteArray<'_, '_> {
2102    type Output = usize;
2103
2104    fn decimal_length(&self, precision: u8, native_length: usize) -> usize {
2105        let length = decimal_length_from_precision(precision);
2106        // The precision of a valid Arrow decimal type fits in its native type
2107        assert!(length <= native_length);
2108        length
2109    }
2110
2111    fn values(self, length: usize, fill: impl Fn(usize, &mut [u8])) -> Result<usize> {
2112        let indices = self.levels.non_null_indices().iter().copied();
2113        let values = chunk_array_slice(length, indices, fill);
2114        self.writer.write_batch_internal(
2115            values.as_slice(),
2116            None,
2117            self.levels.def_level_data().as_ref(),
2118            self.levels.rep_level_data().as_ref(),
2119            None,
2120            None,
2121            None,
2122        )
2123    }
2124}
2125
2126fn write_primitive<E: ColumnValueEncoder>(
2127    writer: &mut GenericColumnWriter<E>,
2128    values: &E::Values,
2129    levels: &ArrayLevels,
2130) -> Result<usize> {
2131    writer.write_batch_internal(
2132        values,
2133        Some(levels.non_null_indices()),
2134        levels.def_level_data().as_ref(),
2135        levels.rep_level_data().as_ref(),
2136        None,
2137        None,
2138        None,
2139    )
2140}
2141
2142fn get_bool_array_slice(
2143    array: &arrow_array::BooleanArray,
2144    indices: impl ExactSizeIterator<Item = usize>,
2145) -> Vec<bool> {
2146    let mut values = Vec::with_capacity(indices.len());
2147    for i in indices {
2148        values.push(array.value(i))
2149    }
2150    values
2151}
2152
2153trait NativeDecimalType: DecimalType {
2154    type NativeBytes: AsRef<[u8]>;
2155
2156    fn to_be_bytes(value: Self::Native) -> Self::NativeBytes;
2157}
2158impl NativeDecimalType for Decimal32Type {
2159    type NativeBytes = [u8; Self::BYTE_LENGTH];
2160
2161    fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
2162        value.to_be_bytes()
2163    }
2164}
2165impl NativeDecimalType for Decimal64Type {
2166    type NativeBytes = [u8; Self::BYTE_LENGTH];
2167
2168    fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
2169        value.to_be_bytes()
2170    }
2171}
2172impl NativeDecimalType for Decimal128Type {
2173    type NativeBytes = [u8; Self::BYTE_LENGTH];
2174
2175    fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
2176        value.to_be_bytes()
2177    }
2178}
2179impl NativeDecimalType for Decimal256Type {
2180    type NativeBytes = [u8; Self::BYTE_LENGTH];
2181
2182    fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
2183        value.to_be_bytes()
2184    }
2185}
2186
2187#[inline]
2188fn chunk_array_slice(
2189    chunk_size: usize,
2190    indices: impl ExactSizeIterator<Item = usize>,
2191    writer: impl Fn(usize, &mut [u8]),
2192) -> Vec<FixedLenByteArray> {
2193    let capacity = indices.len() * chunk_size;
2194    // TODO: This could be done with Vec::spare_capacity_mut,
2195    //       but [MaybeUninit]::write_copy_of_slice is gated behind MSRV 1.93
2196    let mut arena = vec![0; capacity];
2197    for (i, chunk) in indices.zip(arena.chunks_exact_mut(chunk_size)) {
2198        writer(i, chunk);
2199    }
2200    chunk_contiguous_vec(arena, chunk_size)
2201}
2202
2203fn chunk_contiguous_vec(arena: Vec<u8>, chunk_size: usize) -> Vec<FixedLenByteArray> {
2204    let mut values = Vec::with_capacity(arena.len() / chunk_size);
2205    let mut arena = Bytes::from(arena);
2206    while arena.len() >= chunk_size {
2207        let slice = arena.split_to(chunk_size);
2208        values.push(FixedLenByteArray::from(ByteArray::from(slice)));
2209    }
2210    values
2211}
2212
2213/// Hash a byte slice to a u64 for NDV tracking.
2214#[inline]
2215fn hash_bytes(bytes: &[u8]) -> u64 {
2216    twox_hash::XxHash64::oneshot(0, bytes)
2217}
2218
2219/// Returns the fixed byte width for primitive Arrow types, or `None` for variable-length types.
2220fn fixed_byte_width(dt: &ArrowDataType) -> Option<usize> {
2221    use ArrowDataType::*;
2222    match dt {
2223        Int8 | UInt8 => Some(1),
2224        Int16 | UInt16 | Float16 => Some(2),
2225        Int32 | UInt32 | Float32 | Date32 | Time32(_) | Decimal32(_, _) => Some(4),
2226        Int64
2227        | UInt64
2228        | Float64
2229        | Date64
2230        | Time64(_)
2231        | Timestamp(_, _)
2232        | Duration(_)
2233        | Decimal64(_, _) => Some(8),
2234        Interval(IntervalUnit::YearMonth) => Some(4),
2235        Interval(IntervalUnit::DayTime) => Some(8),
2236        Interval(IntervalUnit::MonthDayNano) => Some(16),
2237        Decimal128(_, _) => Some(16),
2238        Decimal256(_, _) => Some(32),
2239        _ => None,
2240    }
2241}
2242
2243/// Hash the non-null values in `array` (at `non_null_indices`) into `seen`.
2244///
2245/// Handles primitive, boolean, fixed-size-binary, and variable-length (Utf8/Binary)
2246/// arrays. Unsupported types are silently skipped, leaving `seen` unchanged for
2247/// those values (NDV is best-effort).
2248fn update_distinct_values_seen(
2249    array: &dyn arrow_array::Array,
2250    non_null_indices: &[usize],
2251    seen: &mut DistinctValuesSet,
2252) {
2253    let data = array.to_data();
2254    let offset = data.offset();
2255
2256    match array.data_type() {
2257        ArrowDataType::Boolean => {
2258            let arr = array
2259                .as_any()
2260                .downcast_ref::<arrow_array::BooleanArray>()
2261                .unwrap();
2262            for &row in non_null_indices {
2263                seen.insert(arr.value(row) as u64);
2264            }
2265        }
2266        ArrowDataType::Utf8 | ArrowDataType::Binary => {
2267            let offsets = data.buffers()[0].typed_data::<i32>();
2268            let values = data.buffers()[1].as_slice();
2269            for &row in non_null_indices {
2270                let start = offsets[offset + row] as usize;
2271                let end = offsets[offset + row + 1] as usize;
2272                seen.insert(hash_bytes(&values[start..end]));
2273            }
2274        }
2275        ArrowDataType::LargeUtf8 | ArrowDataType::LargeBinary => {
2276            let offsets = data.buffers()[0].typed_data::<i64>();
2277            let values = data.buffers()[1].as_slice();
2278            for &row in non_null_indices {
2279                let start = offsets[offset + row] as usize;
2280                let end = offsets[offset + row + 1] as usize;
2281                seen.insert(hash_bytes(&values[start..end]));
2282            }
2283        }
2284        ArrowDataType::FixedSizeBinary(byte_width) => {
2285            let byte_width = *byte_width as usize;
2286            let buffer = data.buffers()[0].as_slice();
2287            for &row in non_null_indices {
2288                let start = (offset + row) * byte_width;
2289                seen.insert(hash_bytes(&buffer[start..start + byte_width]));
2290            }
2291        }
2292        ArrowDataType::Utf8View => {
2293            let string_view_array = array.as_string_view();
2294            for &row in non_null_indices {
2295                seen.insert(hash_bytes(string_view_array.value(row).as_bytes()));
2296            }
2297        }
2298        ArrowDataType::BinaryView => {
2299            let binary_view_array = array.as_binary_view();
2300            for &row in non_null_indices {
2301                seen.insert(hash_bytes(binary_view_array.value(row)));
2302            }
2303        }
2304        data_type => {
2305            if let Some(width) = fixed_byte_width(data_type) {
2306                let buffer = data.buffers()[0].as_slice();
2307                for &row in non_null_indices {
2308                    let pos = (offset + row) * width;
2309                    seen.insert(hash_bytes(&buffer[pos..pos + width]));
2310                }
2311            }
2312            // nested types (List, LargeList, etc.) are Parquet groups, not leaf columns: skip
2313        }
2314    }
2315}
2316
2317// Allow the helpers to use the same imports in unit and integration tests.
2318#[cfg(test)]
2319use crate as parquet_crate;
2320
2321#[cfg(test)]
2322#[path = "../../../tests/arrow_writer/roundtrip_helpers.rs"]
2323mod roundtrip_helpers;
2324
2325#[cfg(test)]
2326mod tests {
2327    use super::roundtrip_helpers::{
2328        RoundTripTest, SMALL_SIZE, required_and_optional, roundtrip_opts,
2329        roundtrip_opts_with_array_validation,
2330    };
2331    use super::*;
2332    use std::cmp::Ordering;
2333    use std::collections::HashMap;
2334
2335    use std::fs::File;
2336
2337    use crate::arrow::ARROW_SCHEMA_META_KEY;
2338    use crate::arrow::arrow_reader::{ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder};
2339    use crate::column::page::{Page, PageReader};
2340    use crate::file::metadata::thrift::PageHeader;
2341    use crate::file::page_index::column_index::ColumnIndexMetaData;
2342    use crate::file::reader::SerializedPageReader;
2343    use crate::parquet_thrift::{ReadThrift, ThriftSliceInputProtocol};
2344    use crate::schema::types::ColumnPath;
2345    use arrow::datatypes::{DataType, Schema};
2346    use arrow::error::Result as ArrowResult;
2347    use arrow::util::data_gen::create_random_array;
2348    use arrow::util::pretty::pretty_format_batches;
2349    use arrow::{array::*, buffer::Buffer, compute::cast};
2350    use arrow_buffer::{IntervalMonthDayNano, NullBuffer, OffsetBuffer};
2351    use arrow_schema::Fields;
2352    use arrow_select::concat::concat;
2353
2354    use crate::basic::{Encoding, EncodingMask};
2355    use crate::data_type::AsBytes;
2356    use crate::file::metadata::{ColumnChunkMetaData, ParquetMetaData, ParquetMetaDataReader};
2357    use crate::file::properties::{
2358        BloomFilterPosition, EnabledStatistics, ReaderProperties, WriterVersion,
2359    };
2360    use crate::file::serialized_reader::ReadOptionsBuilder;
2361    use crate::file::{
2362        reader::{FileReader, SerializedFileReader},
2363        statistics::Statistics,
2364    };
2365
2366    /// A [`PageStore`] that allocates *sparse, non-contiguous* handles and keeps
2367    /// blobs in a `HashMap` — nothing like the default `Vec<Bytes>`. Used to
2368    /// prove the writer relies only on the opaque-handle contract and never on
2369    /// handles being dense `Vec` indices. Records how many blobs were stored.
2370    #[derive(Debug, Default)]
2371    struct RecordingPageStore {
2372        next: u64,
2373        blobs: HashMap<u64, Bytes>,
2374        puts: Arc<std::sync::atomic::AtomicUsize>,
2375    }
2376
2377    impl PageStore for RecordingPageStore {
2378        fn put(&mut self, value: Bytes) -> Result<PageKey> {
2379            // Deliberately non-sequential, never-zero handles.
2380            let id = 100 + self.next * 7;
2381            self.next += 1;
2382            self.puts.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2383            self.blobs.insert(id, value);
2384            Ok(PageKey::new(id))
2385        }
2386
2387        fn take(&mut self, key: PageKey) -> Result<Bytes> {
2388            self.blobs
2389                .remove(&key.get())
2390                .ok_or_else(|| ParquetError::General(format!("missing key {}", key.get())))
2391        }
2392    }
2393
2394    #[derive(Debug)]
2395    struct RecordingPageStoreFactory {
2396        puts: Arc<std::sync::atomic::AtomicUsize>,
2397    }
2398
2399    impl PageStoreFactory for RecordingPageStoreFactory {
2400        fn create(&self, _args: &PageStoreArgs<'_>) -> Result<Box<dyn PageStore>> {
2401            Ok(Box::new(RecordingPageStore {
2402                puts: self.puts.clone(),
2403                ..Default::default()
2404            }))
2405        }
2406    }
2407
2408    /// A custom [`PageStore`] must produce byte-identical files to the in-memory
2409    /// default, across dictionary and non-dictionary columns and multiple row
2410    /// groups (so multiple store instances are exercised).
2411    #[test]
2412    fn custom_page_store_is_byte_identical_to_default() {
2413        let schema = Arc::new(Schema::new(vec![
2414            Field::new("i", DataType::Int32, true),
2415            // A low-cardinality string column to exercise the dictionary path.
2416            Field::new("s", DataType::Utf8, true),
2417        ]));
2418        let i = Int32Array::from(vec![Some(1), None, Some(3), Some(4), Some(5), Some(6)]);
2419        let s = StringArray::from(vec![
2420            Some("a"),
2421            Some("bb"),
2422            Some("a"),
2423            None,
2424            Some("bb"),
2425            Some("ccc"),
2426        ]);
2427        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(i), Arc::new(s)]).unwrap();
2428
2429        // Small row groups so multiple column chunks (hence multiple store
2430        // instances) are produced.
2431        let props = WriterProperties::builder()
2432            .set_max_row_group_row_count(Some(3))
2433            .build();
2434
2435        let write = |factory: Option<Arc<dyn PageStoreFactory>>| {
2436            let mut buffer = Vec::new();
2437            let mut opts = ArrowWriterOptions::new().with_properties(props.clone());
2438            if let Some(factory) = factory {
2439                opts = opts.with_page_store_factory(factory);
2440            }
2441            let mut writer =
2442                ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap();
2443            writer.write(&batch).unwrap();
2444            writer.close().unwrap();
2445            buffer
2446        };
2447
2448        let default_bytes = write(None);
2449
2450        let puts = Arc::new(std::sync::atomic::AtomicUsize::new(0));
2451        let custom_bytes = write(Some(Arc::new(RecordingPageStoreFactory {
2452            puts: puts.clone(),
2453        })));
2454
2455        assert!(
2456            puts.load(std::sync::atomic::Ordering::Relaxed) > 0,
2457            "custom PageStore was never written to"
2458        );
2459        assert_eq!(
2460            default_bytes, custom_bytes,
2461            "a custom PageStore must produce byte-identical output to the default"
2462        );
2463    }
2464
2465    /// A dictionary-encoded column written through the deferred-ordering Arrow
2466    /// path must round-trip correctly even with the offset index disabled, when
2467    /// only the chunk-level dictionary/data page offsets are rewritten (there is
2468    /// no offset index to rebuild). Spans multiple data pages so the
2469    /// dictionary-first reordering is exercised.
2470    #[test]
2471    #[cfg_attr(miri, ignore)] // Takes too long
2472    fn dictionary_column_round_trips_with_offset_index_disabled() {
2473        let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int32, true)]));
2474
2475        // Low cardinality so the column stays dictionary-encoded; enough rows to
2476        // span several data pages within a single row group.
2477        let values: Vec<Option<i32>> = (0..50_000).map(|i| Some(i % 8)).collect();
2478        let array = Int32Array::from(values.clone());
2479        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
2480
2481        let props = WriterProperties::builder()
2482            .set_offset_index_disabled(true)
2483            .set_data_page_row_count_limit(4096)
2484            .build();
2485        let opts = ArrowWriterOptions::new().with_properties(props);
2486
2487        let mut buffer = Vec::new();
2488        let mut writer =
2489            ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap();
2490        writer.write(&batch).unwrap();
2491        writer.close().unwrap();
2492
2493        let reader = ParquetRecordBatchReader::try_new(Bytes::from(buffer), values.len()).unwrap();
2494        let read: Vec<RecordBatch> = reader.collect::<ArrowResult<_>>().unwrap();
2495        let read_values: Vec<Option<i32>> = read
2496            .iter()
2497            .flat_map(|b| b.column(0).as_primitive::<Int32Type>().iter())
2498            .collect();
2499        assert_eq!(read_values, values);
2500    }
2501
2502    /// The dictionary page is routed through the [`PageStore`] like any other
2503    /// page rather than held resident in memory, so a dictionary column chunk's
2504    /// *entire* serialized size — dictionary page included — passes through the
2505    /// store.
2506    #[test]
2507    fn dictionary_page_is_routed_through_the_store() {
2508        /// A store that sums the bytes handed to `put`.
2509        #[derive(Debug, Default)]
2510        struct SizeRecordingPageStore {
2511            blobs: Vec<Bytes>,
2512            bytes_put: Arc<std::sync::atomic::AtomicUsize>,
2513        }
2514        impl PageStore for SizeRecordingPageStore {
2515            fn put(&mut self, value: Bytes) -> Result<PageKey> {
2516                self.bytes_put
2517                    .fetch_add(value.len(), std::sync::atomic::Ordering::Relaxed);
2518                let key = PageKey::new(self.blobs.len() as u64);
2519                self.blobs.push(value);
2520                Ok(key)
2521            }
2522            fn take(&mut self, key: PageKey) -> Result<Bytes> {
2523                Ok(std::mem::take(&mut self.blobs[key.get() as usize]))
2524            }
2525        }
2526        #[derive(Debug)]
2527        struct Factory {
2528            bytes_put: Arc<std::sync::atomic::AtomicUsize>,
2529        }
2530        impl PageStoreFactory for Factory {
2531            fn create(&self, _args: &PageStoreArgs<'_>) -> Result<Box<dyn PageStore>> {
2532                Ok(Box::new(SizeRecordingPageStore {
2533                    bytes_put: self.bytes_put.clone(),
2534                    ..Default::default()
2535                }))
2536            }
2537        }
2538
2539        let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)]));
2540        // Low cardinality keeps the column dictionary-encoded with a real,
2541        // non-empty dictionary page.
2542        let values: Vec<&str> = (0..2048)
2543            .map(|i| ["alpha", "beta", "gamma", "delta"][i % 4])
2544            .collect();
2545        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(StringArray::from(values))])
2546            .unwrap();
2547
2548        let bytes_put = Arc::new(std::sync::atomic::AtomicUsize::new(0));
2549        let opts = ArrowWriterOptions::new().with_page_store_factory(Arc::new(Factory {
2550            bytes_put: bytes_put.clone(),
2551        }));
2552
2553        // A single batch / single column means exactly one row group and one
2554        // store instance, so the bytes it saw map to one column chunk.
2555        let mut buffer = Vec::new();
2556        let mut writer =
2557            ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap();
2558        writer.write(&batch).unwrap();
2559        writer.close().unwrap();
2560
2561        let reader = SerializedFileReader::new(Bytes::from(buffer)).unwrap();
2562        let column = reader.metadata().row_group(0).column(0);
2563        assert!(
2564            column.dictionary_page_offset().is_some(),
2565            "expected the column to be dictionary-encoded"
2566        );
2567
2568        // The bytes the store was handed must account for the whole chunk,
2569        // dictionary page included. Holding the dictionary page apart from the
2570        // store would make this fall short by the dictionary page's size.
2571        assert_eq!(
2572            bytes_put.load(std::sync::atomic::Ordering::Relaxed) as i64,
2573            column.compressed_size(),
2574            "the dictionary page must pass through the store like any other page"
2575        );
2576    }
2577
2578    /// The `page_encoding_stats` list of a dictionary-encoded column chunk must
2579    /// list the dictionary page *before* the data pages, matching the on-disk
2580    /// page layout (the dictionary page is always written first) and the order
2581    /// produced by the column-at-a-time `SerializedFileWriter`.
2582    #[test]
2583    fn dictionary_page_encoding_stats_lists_dictionary_first() {
2584        use crate::basic::PageType;
2585        use crate::file::metadata::ParquetMetaDataOptions;
2586
2587        // A low-cardinality column so it stays dictionary-encoded, with enough
2588        // rows to emit at least one data page in addition to the dictionary page.
2589        let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int32, true)]));
2590        let values: Vec<Option<i32>> = (0..4096).map(|i| Some(i % 4)).collect();
2591        let array = Int32Array::from(values);
2592        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
2593
2594        let mut buffer = Vec::new();
2595        let mut writer = ArrowWriter::try_new(&mut buffer, schema.clone(), None).unwrap();
2596        writer.write(&batch).unwrap();
2597        writer.close().unwrap();
2598        let buffer = Bytes::from(buffer);
2599
2600        // Read the *full*, ordered page encoding stats. By default this crate
2601        // collapses them to a bitmask, which would discard the ordering this
2602        // test is about, so opt out of that here.
2603        let options = ParquetMetaDataOptions::new().with_encoding_stats_as_mask(false);
2604        let metadata = ParquetMetaDataReader::new()
2605            .with_metadata_options(Some(options))
2606            .parse_and_finish(&buffer)
2607            .unwrap();
2608
2609        let col = metadata.row_group(0).column(0);
2610        assert!(
2611            col.dictionary_page_offset().is_some(),
2612            "column should be dictionary encoded"
2613        );
2614        let stats = col
2615            .page_encoding_stats()
2616            .expect("full (non-mask) page encoding stats should be present");
2617
2618        // Arrow C++ (`IsColumnChunkFullyDictionaryEncoded`) requires the
2619        // dictionary page entry to be the first one.
2620        assert_eq!(
2621            stats[0].page_type,
2622            PageType::DICTIONARY_PAGE,
2623            "the DICTIONARY_PAGE encoding stat must come first, got {stats:?}"
2624        );
2625        assert!(
2626            stats[1..]
2627                .iter()
2628                .all(|s| s.page_type == PageType::DATA_PAGE)
2629        );
2630    }
2631
2632    fn get_bytes_after_close(schema: SchemaRef, expected_batch: &RecordBatch) -> Vec<u8> {
2633        let mut buffer = vec![];
2634
2635        let mut writer = ArrowWriter::try_new(&mut buffer, schema, None).unwrap();
2636        writer.write(expected_batch).unwrap();
2637        writer.close().unwrap();
2638
2639        buffer
2640    }
2641
2642    fn get_bytes_by_into_inner(schema: SchemaRef, expected_batch: &RecordBatch) -> Vec<u8> {
2643        let mut writer = ArrowWriter::try_new(Vec::new(), schema, None).unwrap();
2644        writer.write(expected_batch).unwrap();
2645        writer.into_inner().unwrap()
2646    }
2647
2648    /// Writes each of `arrays` as a batch, in several mini-batches, checks they round trip,
2649    /// and returns the number of values in the column's dictionary page
2650    fn write_string_dictionaries(arrays: &[ArrayRef]) -> usize {
2651        let schema = Arc::new(Schema::new(vec![Field::new(
2652            "col",
2653            arrays[0].data_type().clone(),
2654            true,
2655        )]));
2656        let props = WriterProperties::builder().set_write_batch_size(2).build();
2657        let mut file = vec![];
2658        let mut writer = ArrowWriter::try_new(&mut file, schema.clone(), Some(props)).unwrap();
2659        for array in arrays {
2660            let batch = RecordBatch::try_new(schema.clone(), vec![array.clone()]).unwrap();
2661            writer.write(&batch).unwrap();
2662        }
2663        writer.close().unwrap();
2664
2665        let reader = SerializedFileReader::new(Bytes::from(file.clone())).unwrap();
2666        let mut pages = reader
2667            .get_row_group(0)
2668            .unwrap()
2669            .get_column_page_reader(0)
2670            .unwrap();
2671        let dictionary_page = pages.get_next_page().unwrap().unwrap();
2672        assert!(matches!(dictionary_page, Page::DictionaryPage { .. }));
2673
2674        let expected: Vec<&dyn Array> = arrays.iter().map(|a| a.as_ref()).collect();
2675        let expected = concat(&expected).unwrap();
2676        let actual = read_column(file);
2677        assert_eq!(
2678            cast(&actual, &DataType::Utf8).unwrap().as_ref(),
2679            cast(&expected, &DataType::Utf8).unwrap().as_ref()
2680        );
2681
2682        dictionary_page.num_values() as usize
2683    }
2684
2685    #[test]
2686    fn arrow_writer_string_dictionary_cache_reused_for_same_values() {
2687        // "dup" is in the dictionary twice, and values are both short enough to be inlined and not
2688        let shared: ArrayRef = Arc::new(StringArray::from(vec![
2689            "a",
2690            "bb",
2691            "longer than eight bytes",
2692            "dup",
2693            "dup",
2694        ]));
2695        let keys = Int32Array::from(vec![
2696            Some(0),
2697            Some(1),
2698            None,
2699            Some(2),
2700            Some(3),
2701            Some(4),
2702            Some(0),
2703        ]);
2704        let first = DictionaryArray::new(keys, shared.clone());
2705        // Same values `Arc` as `first`, so encoded with the cache built for `first`. Sliced,
2706        // with fewer rows than values, so the cache is only used because it is reused
2707        let keys = Int32Array::from(vec![Some(4), Some(4), Some(1), None, Some(2), Some(0)]);
2708        let second = DictionaryArray::new(keys, shared).slice(1, 4);
2709
2710        let arrays: Vec<ArrayRef> = vec![Arc::new(first), Arc::new(second)];
2711        // a, bb, longer than eight bytes, dup
2712        assert_eq!(write_string_dictionaries(&arrays), 4);
2713    }
2714
2715    #[test]
2716    fn arrow_writer_string_dictionary_cache_replaced_for_new_values() {
2717        let first_values: ArrayRef = Arc::new(StringArray::from(vec!["a", "bb", "ccc"]));
2718        let first = DictionaryArray::new(Int32Array::from(vec![0, 1, 2, 0]), first_values.clone());
2719        // A different values `Arc`, sharing some values with `first` at other keys, so a
2720        // stale cache would map keys to the wrong values
2721        let second_values: ArrayRef = Arc::new(StringArray::from(vec!["zz", "a", "new value!!"]));
2722        let second = DictionaryArray::new(
2723            Int32Array::from(vec![Some(2), Some(1), None, Some(0), Some(1)]),
2724            second_values,
2725        );
2726        // Back to `first`'s values, whose cache has since been replaced
2727        let third = DictionaryArray::new(Int32Array::from(vec![2, 1, 0]), first_values);
2728
2729        let arrays: Vec<ArrayRef> = vec![Arc::new(first), Arc::new(second), Arc::new(third)];
2730        // a, bb, ccc, zz, new value!!
2731        assert_eq!(write_string_dictionaries(&arrays), 5);
2732    }
2733
2734    fn dictionary_writer_memory_size(array: ArrayRef, batch_size: usize) -> usize {
2735        let schema = Arc::new(Schema::new(vec![Field::new(
2736            "col",
2737            array.data_type().clone(),
2738            true,
2739        )]));
2740        let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap();
2741        let props = WriterProperties::builder()
2742            .set_write_batch_size(batch_size)
2743            .build();
2744        let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
2745        writer.write(&batch).unwrap();
2746        writer.memory_size()
2747    }
2748
2749    #[test]
2750    fn arrow_writer_dictionary_cache_ignores_null_rows() {
2751        let values: ArrayRef = Arc::new(StringArray::from(vec![
2752            Some("value"),
2753            None,
2754            Some("unused"),
2755            Some("also unused"),
2756        ]));
2757
2758        for non_null in [0, 2] {
2759            let valid = |i| (3..3 + non_null).contains(&i);
2760            let keys = Int32Array::from_iter((0..10).map(|i| valid(i).then_some(0)));
2761            let null_keys: ArrayRef =
2762                Arc::new(DictionaryArray::new(keys, values.clone()).slice(1, 8));
2763            let keys = Int32Array::from_iter_values((0..10).map(|i| i32::from(!valid(i))));
2764            let null_values: ArrayRef =
2765                Arc::new(DictionaryArray::new(keys, values.clone()).slice(1, 8));
2766            let child: ArrayRef = Arc::new(DictionaryArray::new(
2767                Int32Array::from(vec![0; 10]),
2768                values.clone(),
2769            ));
2770            let parent: ArrayRef = Arc::new(
2771                StructArray::new(
2772                    vec![Arc::new(Field::new(
2773                        "item",
2774                        child.data_type().clone(),
2775                        true,
2776                    ))]
2777                    .into(),
2778                    vec![child],
2779                    Some(NullBuffer::from_iter((0..10).map(valid))),
2780                )
2781                .slice(1, 8),
2782            );
2783
2784            for array in [null_keys, null_values, parent] {
2785                let plain_type = match array.data_type() {
2786                    DataType::Struct(_) => DataType::Struct(
2787                        vec![Arc::new(Field::new("item", DataType::Utf8, true))].into(),
2788                    ),
2789                    _ => DataType::Utf8,
2790                };
2791                let plain = cast(&array, &plain_type).unwrap();
2792                // Neither input should retain a cache for only zero or two encoded values.
2793                assert_eq!(
2794                    dictionary_writer_memory_size(array, 2),
2795                    dictionary_writer_memory_size(plain, 2)
2796                );
2797            }
2798        }
2799    }
2800
2801    #[test]
2802    fn arrow_writer_dictionary_cache_counts_rows_across_mini_batches() {
2803        let values: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c", "d"]));
2804        let array: ArrayRef = Arc::new(DictionaryArray::new(
2805            Int32Array::from(vec![
2806                Some(0),
2807                Some(1),
2808                None,
2809                Some(2),
2810                Some(3),
2811                Some(0),
2812                Some(1),
2813                Some(2),
2814                Some(3),
2815            ]),
2816            values,
2817        ));
2818        let plain = cast(&array, &DataType::Utf8).unwrap();
2819        // Eight rows justify a four-slot cache, even with only two rows per mini-batch.
2820        assert_eq!(
2821            dictionary_writer_memory_size(array, 2),
2822            dictionary_writer_memory_size(plain, 2) + 4 * std::mem::size_of::<u64>()
2823        );
2824    }
2825
2826    #[test]
2827    fn roundtrip_bytes() {
2828        // define schema
2829        let schema = Arc::new(Schema::new(vec![
2830            Field::new("a", DataType::Int32, false),
2831            Field::new("b", DataType::Int32, true),
2832        ]));
2833
2834        // create some data
2835        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2836        let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2837
2838        // build a record batch
2839        let expected_batch =
2840            RecordBatch::try_new(schema.clone(), vec![Arc::new(a), Arc::new(b)]).unwrap();
2841
2842        for buffer in [
2843            get_bytes_after_close(schema.clone(), &expected_batch),
2844            get_bytes_by_into_inner(schema, &expected_batch),
2845        ] {
2846            let cursor = Bytes::from(buffer);
2847            let mut record_batch_reader = ParquetRecordBatchReader::try_new(cursor, 1024).unwrap();
2848
2849            let actual_batch = record_batch_reader
2850                .next()
2851                .expect("No batch found")
2852                .expect("Unable to get batch");
2853
2854            assert_eq!(expected_batch.schema(), actual_batch.schema());
2855            assert_eq!(expected_batch.num_columns(), actual_batch.num_columns());
2856            assert_eq!(expected_batch.num_rows(), actual_batch.num_rows());
2857            for i in 0..expected_batch.num_columns() {
2858                let expected_data = expected_batch.column(i).to_data();
2859                let actual_data = actual_batch.column(i).to_data();
2860
2861                assert_eq!(expected_data, actual_data);
2862            }
2863        }
2864    }
2865
2866    fn read_column(file: Vec<u8>) -> ArrayRef {
2867        let reader = ParquetRecordBatchReader::try_new(Bytes::from(file), 4096).unwrap();
2868        let batches: Vec<RecordBatch> = reader.map(|b| b.unwrap()).collect();
2869        let arrays: Vec<&dyn Array> = batches.iter().map(|b| b.column(0).as_ref()).collect();
2870        concat(&arrays).unwrap()
2871    }
2872
2873    fn roundtrip_compatible_column(field: Field, col: ArrayRef) -> ArrayRef {
2874        let writer_schema = Arc::new(Schema::new(vec![field]));
2875        let batch_schema = Arc::new(Schema::new(vec![Field::new(
2876            "c",
2877            col.data_type().clone(),
2878            col.logical_null_count() != 0,
2879        )]));
2880        let batch = RecordBatch::try_new(batch_schema, vec![col]).unwrap();
2881        // Isolate level/value compatibility from Arrow metadata reconstruction,
2882        // which is covered separately by the default-reader regression below.
2883        let options = ArrowWriterOptions::new().with_skip_arrow_metadata(true);
2884        let mut file = vec![];
2885        let mut writer =
2886            ArrowWriter::try_new_with_options(&mut file, writer_schema, options).unwrap();
2887        writer.write(&batch).unwrap();
2888        writer.close().unwrap();
2889        read_column(file)
2890    }
2891
2892    fn assert_compatible_roundtrip(field: Field, actual: ArrayRef, expected: &ArrayRef) {
2893        assert!(!compute_leaves(&field, &actual).unwrap().is_empty());
2894        assert_eq!(
2895            roundtrip_compatible_column(field, actual).as_ref(),
2896            expected.as_ref()
2897        );
2898    }
2899
2900    fn run_end_type(value_type: DataType, nullable_values: bool) -> DataType {
2901        DataType::RunEndEncoded(
2902            Arc::new(Field::new(
2903                Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME,
2904                DataType::Int32,
2905                false,
2906            )),
2907            Arc::new(Field::new(
2908                Field::REE_VALUES_FIELD_DEFAULT_NAME,
2909                value_type,
2910                nullable_values,
2911            )),
2912        )
2913    }
2914
2915    #[test]
2916    fn arrow_writer_required_ree_schema_rejects_nullable_dense_input() {
2917        let writer_schema = Arc::new(Schema::new(vec![Field::new(
2918            "c",
2919            run_end_type(DataType::Int32, false),
2920            false,
2921        )]));
2922        let batch_schema = Arc::new(Schema::new(vec![Field::new("c", DataType::Int32, true)]));
2923        let batch = RecordBatch::try_new(
2924            batch_schema.clone(),
2925            vec![Arc::new(Int32Array::from(vec![Some(7), None, Some(9)]))],
2926        )
2927        .unwrap();
2928
2929        for skip_metadata in [false, true] {
2930            for cdc in [false, true] {
2931                let props = WriterProperties::builder()
2932                    .set_content_defined_chunking(cdc.then(Default::default))
2933                    .build();
2934                let options = ArrowWriterOptions::new()
2935                    .with_properties(props)
2936                    .with_skip_arrow_metadata(skip_metadata);
2937                let mut writer = ArrowWriter::try_new_with_options(
2938                    Vec::new(),
2939                    writer_schema.clone(),
2940                    options.clone(),
2941                )
2942                .unwrap();
2943                let error = writer.write(&batch).unwrap_err().to_string();
2944                assert!(
2945                    error.contains("Found null at index 1 for required field 'c'"),
2946                    "skip_metadata={skip_metadata}, cdc={cdc}: {error}"
2947                );
2948
2949                // A nullable batch schema is fine when the data has no actual nulls.
2950                let values: ArrayRef = Arc::new(Int32Array::from(vec![7, 8, 9]));
2951                let batch =
2952                    RecordBatch::try_new(batch_schema.clone(), vec![values.clone()]).unwrap();
2953                let mut file = vec![];
2954                let mut writer =
2955                    ArrowWriter::try_new_with_options(&mut file, writer_schema.clone(), options)
2956                        .unwrap();
2957                writer.write(&batch).unwrap();
2958                writer.close().unwrap();
2959                assert_eq!(read_column(file).as_ref(), values.as_ref());
2960            }
2961        }
2962    }
2963
2964    #[test]
2965    fn ree_nulls_respect_dense_schema_nullability() {
2966        let run_ends = Int32Array::from(vec![2, 4, 5]);
2967        let values = Int32Array::from(vec![Some(1), None, Some(2)]);
2968        let ree: ArrayRef = Arc::new(Int32RunArray::try_new(&run_ends, &values).unwrap());
2969        let ree_batch = RecordBatch::try_new(
2970            Arc::new(Schema::new(vec![Field::new(
2971                "c",
2972                ree.data_type().clone(),
2973                true,
2974            )])),
2975            vec![ree],
2976        )
2977        .unwrap();
2978
2979        let nullable_schema = Arc::new(Schema::new(vec![Field::new("c", DataType::Int32, true)]));
2980        let mut file = vec![];
2981        let mut writer = ArrowWriter::try_new(&mut file, nullable_schema, None).unwrap();
2982        writer.write(&ree_batch).unwrap();
2983        writer.close().unwrap();
2984        assert_eq!(
2985            read_column(file).as_primitive::<Int32Type>(),
2986            &Int32Array::from(vec![Some(1), Some(1), None, None, Some(2)])
2987        );
2988
2989        let required_schema = Arc::new(Schema::new(vec![Field::new("c", DataType::Int32, false)]));
2990        let mut writer = ArrowWriter::try_new(Vec::new(), required_schema, None).unwrap();
2991        let error = writer.write(&ree_batch).unwrap_err().to_string();
2992        assert!(error.contains("Found null") && error.contains("required field 'c'"));
2993    }
2994
2995    #[test]
2996    fn required_null_validation_respects_reachability_and_write_timing() {
2997        let required = Field::new("c", DataType::Int32, false);
2998        let nullable: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), None]));
2999        let leaves = compute_leaves(&required, &nullable).unwrap();
3000        let schema = Arc::new(Schema::new(vec![required.clone()]));
3001        let writer = ArrowWriter::try_new(Vec::new(), schema, None).unwrap();
3002        let (_file, factory) = writer.into_serialized_writer().unwrap();
3003        let mut columns = factory.create_column_writers(0).unwrap();
3004        assert!(
3005            columns[0]
3006                .write(&leaves[0])
3007                .unwrap_err()
3008                .to_string()
3009                .contains("required field 'c'")
3010        );
3011
3012        // Only a referenced logical dictionary null is invalid, not an unused value.
3013        let values: ArrayRef = Arc::new(Int32Array::from(vec![Some(7), None]));
3014        let dict: ArrayRef = Arc::new(DictionaryArray::new(
3015            Int8Array::from(vec![0, 0]),
3016            values.clone(),
3017        ));
3018        assert_eq!(
3019            roundtrip_compatible_column(required.clone(), dict).as_primitive::<Int32Type>(),
3020            &Int32Array::from(vec![7, 7])
3021        );
3022        let dict: ArrayRef = Arc::new(DictionaryArray::new(Int8Array::from(vec![0, 1]), values));
3023        let leaves = compute_leaves(&required, &dict).unwrap();
3024        assert!(leaves[0].0.validate().is_err());
3025        let dict: ArrayRef = Arc::new(DictionaryArray::new(
3026            Int8Array::from(vec![Some(0), None]),
3027            Arc::new(Int32Array::from(vec![7])),
3028        ));
3029        let leaves = compute_leaves(&required, &dict).unwrap();
3030        assert!(leaves[0].0.validate().is_err());
3031
3032        // A nullable ancestor masks an otherwise required child null, including
3033        // when the target child has a schema-only REE wrapper.
3034        let actual_fields = Fields::from(vec![Field::new("source_child", DataType::Int32, true)]);
3035        let array: ArrayRef = Arc::new(StructArray::new(
3036            actual_fields.clone(),
3037            vec![nullable.clone()],
3038            Some(NullBuffer::from(vec![true, false])),
3039        ));
3040        let unmasked: ArrayRef = Arc::new(StructArray::new(
3041            actual_fields,
3042            vec![nullable.clone()],
3043            None,
3044        ));
3045        for child_type in [DataType::Int32, run_end_type(DataType::Int32, false)] {
3046            let target = Field::new(
3047                "c",
3048                DataType::Struct(Fields::from(vec![Field::new(
3049                    "target_child",
3050                    child_type,
3051                    false,
3052                )])),
3053                true,
3054            );
3055            let leaves = compute_leaves(&target, &unmasked).unwrap();
3056            assert!(leaves[0].0.validate().is_err());
3057            let leaves = compute_leaves(&target, &array).unwrap();
3058            assert!(leaves[0].0.validate().is_ok());
3059            let result = roundtrip_compatible_column(target, array.clone());
3060            assert_eq!(result.as_struct().fields()[0].name(), "target_child");
3061            assert!(result.is_null(1));
3062        }
3063
3064        // Unused list child slots and out-of-slice scalar nulls are unreachable.
3065        let values: ArrayRef = Arc::new(Int32Array::from(vec![None, Some(9), None]));
3066        let list: ArrayRef = Arc::new(ListArray::new(
3067            Arc::new(Field::new_list_field(DataType::Int32, true)),
3068            OffsetBuffer::new(vec![1_i32, 2].into()),
3069            values,
3070            None,
3071        ));
3072        let target = Field::new(
3073            "c",
3074            DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false))),
3075            false,
3076        );
3077        assert!(
3078            compute_leaves(&target, &list).unwrap()[0]
3079                .0
3080                .validate()
3081                .is_ok()
3082        );
3083        assert_eq!(
3084            roundtrip_compatible_column(required, nullable.slice(0, 1)).as_primitive::<Int32Type>(),
3085            &Int32Array::from(vec![1])
3086        );
3087    }
3088
3089    fn check_required_list_validation(list_type: fn(Arc<Field>) -> DataType) {
3090        let make_array = |values: Vec<Option<i32>>,
3091                          offsets: Vec<i32>,
3092                          sizes: Vec<i32>,
3093                          nulls: Option<NullBuffer>| {
3094            let child = Arc::new(Field::new("source_child", DataType::Int32, true));
3095            let values: ArrayRef = Arc::new(Int32Array::from(values));
3096            match list_type(child.clone()) {
3097                DataType::List(_) => Arc::new(ListArray::new(
3098                    child,
3099                    OffsetBuffer::new(offsets.into()),
3100                    values,
3101                    nulls,
3102                )) as ArrayRef,
3103                DataType::LargeList(_) => Arc::new(LargeListArray::new(
3104                    child,
3105                    OffsetBuffer::new(
3106                        offsets
3107                            .into_iter()
3108                            .map(i64::from)
3109                            .collect::<Vec<_>>()
3110                            .into(),
3111                    ),
3112                    values,
3113                    nulls,
3114                )) as ArrayRef,
3115                DataType::ListView(_) => Arc::new(ListViewArray::new(
3116                    child,
3117                    offsets.into(),
3118                    sizes.into(),
3119                    values,
3120                    nulls,
3121                )) as ArrayRef,
3122                DataType::LargeListView(_) => Arc::new(LargeListViewArray::new(
3123                    child,
3124                    offsets
3125                        .into_iter()
3126                        .map(i64::from)
3127                        .collect::<Vec<_>>()
3128                        .into(),
3129                    sizes.into_iter().map(i64::from).collect::<Vec<_>>().into(),
3130                    values,
3131                    nulls,
3132                )) as ArrayRef,
3133                _ => unreachable!(),
3134            }
3135        };
3136        let is_view = matches!(
3137            list_type(Arc::new(Field::new("item", DataType::Int32, true))),
3138            DataType::ListView(_) | DataType::LargeListView(_)
3139        );
3140        let offsets = |starts: Vec<i32>, end: i32| {
3141            let mut result = starts;
3142            if !is_view {
3143                result.push(end);
3144            }
3145            result
3146        };
3147        let target = Field::new(
3148            "c",
3149            list_type(Arc::new(Field::new("target_child", DataType::Int32, false))),
3150            true,
3151        );
3152        let invalid = make_array(vec![None, Some(7)], offsets(vec![0], 2), vec![2], None);
3153        let valid = make_array(vec![Some(7), Some(8)], offsets(vec![0], 2), vec![2], None);
3154        let masked = make_array(
3155            vec![None, Some(7)],
3156            offsets(vec![0, 1], 2),
3157            vec![1, 1],
3158            Some(NullBuffer::from(vec![false, true])),
3159        );
3160        let unused = make_array(
3161            vec![None, Some(7), None],
3162            offsets(vec![1], 2),
3163            vec![1],
3164            None,
3165        );
3166        let empty = make_array(vec![None], offsets(vec![0], 0), vec![0], None);
3167        let sliced = make_array(
3168            vec![None, Some(7)],
3169            offsets(vec![0, 1], 2),
3170            vec![1, 1],
3171            None,
3172        )
3173        .slice(1, 1);
3174        let schema = Arc::new(Schema::new(vec![target.clone()]));
3175
3176        // Binding succeeds even for invalid input; writing the affected leaf rejects it.
3177        for (name, array, invalid) in [
3178            ("reachable null", invalid, true),
3179            ("valid required child", valid.clone(), false),
3180            ("null ancestor", masked, false),
3181            ("unused child slots", unused, false),
3182            ("empty list", empty, false),
3183            ("out-of-slice null", sliced, false),
3184        ] {
3185            let leaves = compute_leaves(&target, &array).unwrap();
3186            assert_eq!(leaves.len(), 1);
3187            for cdc in [false, true] {
3188                let props = WriterProperties::builder()
3189                    .set_content_defined_chunking(cdc.then(Default::default))
3190                    .build();
3191                let mut writer =
3192                    ArrowWriter::try_new(Vec::new(), schema.clone(), Some(props)).unwrap();
3193                let mut chunkers = writer.cdc_chunkers.take();
3194                let (_file, factory) = writer.into_serialized_writer().unwrap();
3195                let mut columns = factory.create_column_writers(0).unwrap();
3196                let before = (
3197                    columns[0].memory_size(),
3198                    columns[0].get_estimated_total_bytes(),
3199                );
3200                let result = match &mut chunkers {
3201                    Some(chunkers) => columns[0].write_with_chunker(&leaves[0], &mut chunkers[0]),
3202                    None => columns[0].write(&leaves[0]),
3203                };
3204                if invalid {
3205                    let error = result.unwrap_err().to_string();
3206                    assert!(
3207                        error.contains("Found null at index 0 for required field 'target_child'"),
3208                        "{name}, cdc={cdc}: {error}"
3209                    );
3210                    assert_eq!(
3211                        before,
3212                        (
3213                            columns[0].memory_size(),
3214                            columns[0].get_estimated_total_bytes()
3215                        ),
3216                        "invalid leaf must be rejected before encoding"
3217                    );
3218                    // A rejected leaf write must leave no encoded values behind.
3219                    let valid_leaves = compute_leaves(&target, &valid).unwrap();
3220                    match &mut chunkers {
3221                        Some(chunkers) => columns[0]
3222                            .write_with_chunker(&valid_leaves[0], &mut chunkers[0])
3223                            .unwrap(),
3224                        None => columns[0].write(&valid_leaves[0]).unwrap(),
3225                    }
3226                    assert_eq!(
3227                        columns
3228                            .pop()
3229                            .unwrap()
3230                            .close()
3231                            .unwrap()
3232                            .close
3233                            .metadata
3234                            .num_values(),
3235                        2
3236                    );
3237                } else {
3238                    result.unwrap();
3239                    columns.pop().unwrap().close().unwrap();
3240                }
3241
3242                // Exercise the ordinary and CDC row-group driving entry points as well.
3243                let batch = RecordBatch::try_new(
3244                    Arc::new(Schema::new(vec![Field::new(
3245                        "c",
3246                        array.data_type().clone(),
3247                        true,
3248                    )])),
3249                    vec![array.clone()],
3250                )
3251                .unwrap();
3252                let props = WriterProperties::builder()
3253                    .set_content_defined_chunking(cdc.then(Default::default))
3254                    .build();
3255                let mut writer =
3256                    ArrowWriter::try_new(Vec::new(), schema.clone(), Some(props)).unwrap();
3257                let result = writer.write(&batch);
3258                if invalid {
3259                    assert!(
3260                        result
3261                            .unwrap_err()
3262                            .to_string()
3263                            .contains("required field 'target_child'")
3264                    );
3265                } else {
3266                    result.unwrap();
3267                    writer.close().unwrap();
3268                }
3269            }
3270        }
3271    }
3272
3273    #[test]
3274    fn required_list_validation_preserves_level_generation() {
3275        check_required_list_validation(DataType::List);
3276    }
3277
3278    #[test]
3279    fn required_large_list_validation_preserves_level_generation() {
3280        check_required_list_validation(DataType::LargeList);
3281    }
3282
3283    #[test]
3284    fn required_list_view_validation_preserves_level_generation() {
3285        check_required_list_validation(DataType::ListView);
3286    }
3287
3288    #[test]
3289    fn required_large_list_view_validation_preserves_level_generation() {
3290        check_required_list_validation(DataType::LargeListView);
3291    }
3292
3293    fn nullable_int32_list() -> ArrayRef {
3294        Arc::new(ListArray::new(
3295            Arc::new(Field::new_list_field(DataType::Int32, true)),
3296            OffsetBuffer::new(vec![0_i32, 2, 3].into()),
3297            Arc::new(Int32Array::from(vec![Some(1), None, Some(3)])),
3298            None,
3299        ))
3300    }
3301
3302    #[test]
3303    fn arrow_writer_dense_list_under_ree_schema() {
3304        let list = nullable_int32_list();
3305        // The schema wrapper is physical, not a second logical list node.
3306        let field = Field::new("c", run_end_type(list.data_type().clone(), false), false);
3307        assert_compatible_roundtrip(field, list.clone(), &list);
3308    }
3309
3310    #[test]
3311    fn arrow_writer_dense_list_under_dictionary_schema() {
3312        let list = nullable_int32_list();
3313        let field = Field::new(
3314            "c",
3315            DataType::Dictionary(Box::new(DataType::Int8), Box::new(list.data_type().clone())),
3316            false,
3317        );
3318        assert_compatible_roundtrip(field, list.clone(), &list);
3319    }
3320
3321    #[test]
3322    fn arrow_writer_dense_values_under_nested_ree_schema() {
3323        let dense: ArrayRef = Arc::new(Int32Array::from(vec![Some(7), None, Some(9)]));
3324        let field = Field::new(
3325            "c",
3326            run_end_type(run_end_type(DataType::Int32, true), false),
3327            false,
3328        );
3329        assert_compatible_roundtrip(field, dense.clone(), &dense);
3330    }
3331
3332    #[test]
3333    fn arrow_writer_dense_values_under_dictionary_ree_schema() {
3334        let dense: ArrayRef = Arc::new(Int32Array::from(vec![Some(7), None, Some(9)]));
3335        let field = Field::new(
3336            "c",
3337            DataType::Dictionary(
3338                Box::new(DataType::Int8),
3339                Box::new(run_end_type(DataType::Int32, true)),
3340            ),
3341            false,
3342        );
3343        assert_compatible_roundtrip(field, dense.clone(), &dense);
3344    }
3345
3346    #[test]
3347    fn arrow_writer_dense_values_under_ree_dictionary_schema() {
3348        let dense: ArrayRef = Arc::new(Int32Array::from(vec![Some(7), None, Some(9)]));
3349        // Hoist REE value nullability before peeling the dictionary; the
3350        // dense batch contains an actual null despite the required outer field.
3351        let field = Field::new(
3352            "c",
3353            run_end_type(
3354                DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int32)),
3355                true,
3356            ),
3357            false,
3358        );
3359        assert_compatible_roundtrip(field, dense.clone(), &dense);
3360    }
3361
3362    #[test]
3363    fn arrow_writer_dense_struct_under_ree_child_schema() {
3364        let dense: ArrayRef = Arc::new(Int32Array::from(vec![Some(7), None, Some(9)]));
3365        let struct_fields = Fields::from(vec![Field::new("value", DataType::Int32, true)]);
3366        let dense_struct: ArrayRef = Arc::new(StructArray::new(struct_fields, vec![dense], None));
3367        // A wrapper on a child must not make the enclosing struct incompatible.
3368        let field = Field::new(
3369            "c",
3370            DataType::Struct(Fields::from(vec![Field::new(
3371                "value",
3372                run_end_type(DataType::Int32, true),
3373                false,
3374            )])),
3375            false,
3376        );
3377        assert_compatible_roundtrip(field, dense_struct.clone(), &dense_struct);
3378    }
3379
3380    #[test]
3381    fn arrow_writer_dense_list_under_ree_child_schema() {
3382        let list = nullable_int32_list();
3383        let field = Field::new(
3384            "c",
3385            DataType::List(Arc::new(Field::new(
3386                "item",
3387                run_end_type(DataType::Int32, true),
3388                false,
3389            ))),
3390            false,
3391        );
3392        assert_compatible_roundtrip(field, list.clone(), &list);
3393    }
3394
3395    #[test]
3396    fn arrow_writer_dense_list_under_dictionary_child_schema() {
3397        let list = nullable_int32_list();
3398        let field = Field::new(
3399            "c",
3400            DataType::List(Arc::new(Field::new_list_field(
3401                DataType::Dictionary(Box::new(DataType::Int8), Box::new(DataType::Int32)),
3402                true,
3403            ))),
3404            false,
3405        );
3406        assert_compatible_roundtrip(field, list.clone(), &list);
3407    }
3408
3409    #[test]
3410    fn arrow_writer_dictionary_list_under_dense_child_schema() {
3411        let dict: ArrayRef = Arc::new(DictionaryArray::new(
3412            Int8Array::from(vec![Some(0), None, Some(1)]),
3413            Arc::new(Int32Array::from(vec![1, 3])),
3414        ));
3415        let dict_list: ArrayRef = Arc::new(ListArray::new(
3416            Arc::new(Field::new_list_field(dict.data_type().clone(), true)),
3417            OffsetBuffer::new(vec![0_i32, 2, 3].into()),
3418            dict,
3419            None,
3420        ));
3421        let expected = nullable_int32_list();
3422        assert_compatible_roundtrip(
3423            Field::new("c", expected.data_type().clone(), false),
3424            dict_list,
3425            &expected,
3426        );
3427    }
3428
3429    #[test]
3430    fn arrow_writer_dense_map_under_ree_value_schema() {
3431        let key_field = Arc::new(Field::new("keys", DataType::Utf8, false));
3432        let entries = StructArray::new(
3433            Fields::from(vec![
3434                key_field.clone(),
3435                Arc::new(Field::new("values", DataType::Int32, true)),
3436            ]),
3437            vec![
3438                Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef,
3439                Arc::new(Int32Array::from(vec![Some(7), None, Some(9)])),
3440            ],
3441            None,
3442        );
3443        let map: ArrayRef = Arc::new(MapArray::new(
3444            Arc::new(Field::new("entries", entries.data_type().clone(), false)),
3445            OffsetBuffer::new(vec![0_i32, 2, 3].into()),
3446            entries,
3447            None,
3448            false,
3449        ));
3450        // Resolve the value wrapper after walking through the entries struct.
3451        let field = Field::new(
3452            "c",
3453            DataType::Map(
3454                Arc::new(Field::new(
3455                    "entries",
3456                    DataType::Struct(Fields::from(vec![
3457                        key_field,
3458                        Arc::new(Field::new(
3459                            "values",
3460                            run_end_type(DataType::Int32, true),
3461                            false,
3462                        )),
3463                    ])),
3464                    false,
3465                )),
3466                false,
3467            ),
3468            false,
3469        );
3470        assert_compatible_roundtrip(field, map.clone(), &map);
3471    }
3472
3473    #[test]
3474    fn arrow_writer_nested_dictionary_metadata_roundtrips_with_default_reader() {
3475        let dense: ArrayRef = Arc::new(StructArray::new(
3476            Fields::from(vec![Field::new("item", DataType::Int32, true)]),
3477            vec![Arc::new(Int32Array::from(vec![Some(7), None, Some(9)]))],
3478            None,
3479        ));
3480        let field = Field::new(
3481            "value",
3482            DataType::Dictionary(
3483                Box::new(DataType::Int8),
3484                Box::new(dense.data_type().clone()),
3485            ),
3486            false,
3487        )
3488        .with_metadata(HashMap::from([(
3489            "description".to_string(),
3490            "nested dictionary".to_string(),
3491        )]));
3492        let writer_schema = Arc::new(Schema::new(vec![field.clone()]));
3493        let expected_schema = Arc::new(Schema::new(vec![
3494            field.with_data_type(dense.data_type().clone()),
3495        ]));
3496        let batch = RecordBatch::try_new(expected_schema.clone(), vec![dense]).unwrap();
3497
3498        // Keep Arrow metadata enabled: ignoring it would hide the original
3499        // reader-initialization failure caused by the nested dictionary hint.
3500        let mut file = vec![];
3501        let mut writer = ArrowWriter::try_new(&mut file, writer_schema, None).unwrap();
3502        writer.write(&batch).unwrap();
3503        writer.close().unwrap();
3504
3505        let builder = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(file)).unwrap();
3506        assert!(
3507            builder
3508                .metadata()
3509                .file_metadata()
3510                .key_value_metadata()
3511                .unwrap()
3512                .iter()
3513                .any(|kv| kv.key == ARROW_SCHEMA_META_KEY)
3514        );
3515        assert_eq!(builder.schema(), &expected_schema);
3516        let decoded = builder
3517            .build()
3518            .unwrap()
3519            .collect::<ArrowResult<Vec<_>>>()
3520            .unwrap();
3521        assert_eq!(decoded, vec![batch]);
3522    }
3523
3524    #[test]
3525    fn arrow_writer_rejects_missing_or_extra_top_level_columns_cleanly() {
3526        let one_field = Arc::new(Schema::new(vec![Field::new(
3527            "a",
3528            ArrowDataType::Int32,
3529            false,
3530        )]));
3531        let two_fields = Arc::new(Schema::new(vec![
3532            Field::new("a", ArrowDataType::Int32, false),
3533            Field::new("b", ArrowDataType::Int32, false),
3534        ]));
3535        let one_column = RecordBatch::try_new(
3536            Arc::clone(&one_field),
3537            vec![Arc::new(Int32Array::from(vec![1]))],
3538        )
3539        .unwrap();
3540        let two_columns = RecordBatch::try_new(
3541            Arc::clone(&two_fields),
3542            vec![
3543                Arc::new(Int32Array::from(vec![1])),
3544                Arc::new(Int32Array::from(vec![2])),
3545            ],
3546        )
3547        .unwrap();
3548
3549        let mut writer = ArrowWriter::try_new(Vec::new(), Arc::clone(&two_fields), None).unwrap();
3550        let err = writer.write(&one_column).unwrap_err();
3551        assert!(
3552            err.to_string()
3553                .contains("writer has 2 top-level fields but batch has 1 columns"),
3554            "{err}"
3555        );
3556        assert_eq!(writer.in_progress_rows(), 0);
3557
3558        let mut writer = ArrowWriter::try_new(Vec::new(), Arc::clone(&one_field), None).unwrap();
3559        let err = writer.write(&two_columns).unwrap_err();
3560        assert!(
3561            err.to_string()
3562                .contains("writer has 1 top-level fields but batch has 2 columns"),
3563            "{err}"
3564        );
3565        assert_eq!(writer.in_progress_rows(), 0);
3566    }
3567
3568    #[test]
3569    fn arrow_row_group_factory_rejects_parquet_leaf_count_mismatch_cleanly() {
3570        let one_field = Arc::new(Schema::new(vec![Field::new(
3571            "a",
3572            ArrowDataType::Int32,
3573            false,
3574        )]));
3575        let two_fields = Arc::new(Schema::new(vec![
3576            Field::new("a", ArrowDataType::Int32, false),
3577            Field::new("b", ArrowDataType::Int32, false),
3578        ]));
3579        let parquet = ArrowSchemaConverter::new().convert(&one_field).unwrap();
3580        let props = Arc::new(WriterProperties::default());
3581        let file_writer =
3582            SerializedFileWriter::new(Vec::new(), parquet.root_schema_ptr(), Arc::clone(&props))
3583                .unwrap();
3584
3585        let eager_error =
3586            ArrowRowGroupWriterFactory::try_new(&file_writer, Arc::clone(&two_fields)).unwrap_err();
3587        assert!(matches!(&eager_error, ParquetError::General(_)));
3588        #[expect(deprecated)] // Verify that the legacy constructor still defers errors.
3589        let factory = ArrowRowGroupWriterFactory::new(&file_writer, Arc::clone(&two_fields));
3590        let Err(err) = factory.create_column_writers(0) else {
3591            panic!("mismatched schemas should not create column writers")
3592        };
3593        assert_eq!(err.to_string(), eager_error.to_string());
3594        assert!(
3595            err.to_string()
3596                .contains("Arrow schema maps to 2 leaf columns but Parquet schema has 1"),
3597            "{err}"
3598        );
3599
3600        let parquet = ArrowSchemaConverter::new().convert(&two_fields).unwrap();
3601        let file_writer =
3602            SerializedFileWriter::new(Vec::new(), parquet.root_schema_ptr(), props).unwrap();
3603        let eager_error =
3604            ArrowRowGroupWriterFactory::try_new(&file_writer, Arc::clone(&one_field)).unwrap_err();
3605        assert!(matches!(&eager_error, ParquetError::General(_)));
3606        #[expect(deprecated)] // Verify that the legacy constructor still defers errors.
3607        let factory = ArrowRowGroupWriterFactory::new(&file_writer, one_field);
3608        let Err(err) = factory.create_column_writers(0) else {
3609            panic!("mismatched schemas should not create column writers")
3610        };
3611        assert_eq!(err.to_string(), eager_error.to_string());
3612        assert!(
3613            err.to_string()
3614                .contains("Arrow schema maps to 1 leaf columns but Parquet schema has 2"),
3615            "{err}"
3616        );
3617    }
3618
3619    #[test]
3620    fn arrow_writer_required_map_keys_reject_only_reachable_nulls() {
3621        let keys: ArrayRef = Arc::new(DictionaryArray::<Int32Type>::new(
3622            Int32Array::from(vec![1, 0, 1]),
3623            Arc::new(Int32Array::from(vec![None, Some(7)])),
3624        ));
3625        let values = Int32Array::from(vec![Some(10), Some(20), None]);
3626        let fields = Fields::from(vec![
3627            Field::new("key", keys.data_type().clone(), false),
3628            Field::new("value", DataType::Int32, true),
3629        ]);
3630        assert!(
3631            StructArray::try_new(
3632                fields.clone(),
3633                vec![keys.clone(), Arc::new(values.clone())],
3634                None
3635            )
3636            .is_err()
3637        );
3638
3639        // ArrayData validates physical nulls, not null dictionary values. This
3640        // safe construction reaches the writer with a logically null map key.
3641        let data = ArrayData::builder(DataType::Struct(fields))
3642            .len(3)
3643            .add_child_data(keys.to_data())
3644            .add_child_data(values.to_data())
3645            .build()
3646            .unwrap();
3647        data.validate_full().unwrap();
3648        let entries = StructArray::from(data);
3649        let entries_field = Arc::new(Field::new("entries", entries.data_type().clone(), false));
3650        let offsets = OffsetBuffer::new(vec![0i32, 1, 2, 3].into());
3651        let map = MapArray::try_new(
3652            entries_field.clone(),
3653            offsets.clone(),
3654            entries.clone(),
3655            None,
3656            false,
3657        )
3658        .unwrap();
3659        assert_eq!(map.keys().logical_null_count(), 1);
3660        let masked = MapArray::try_new(
3661            entries_field,
3662            offsets,
3663            entries,
3664            Some(NullBuffer::from(vec![true, false, true])),
3665            false,
3666        )
3667        .unwrap();
3668        let writer_schema = Arc::new(Schema::new(vec![Field::new(
3669            "map",
3670            DataType::Map(
3671                Arc::new(Field::new(
3672                    "entries",
3673                    DataType::Struct(
3674                        vec![
3675                            Field::new("key", DataType::Int32, false),
3676                            Field::new("value", DataType::Int32, true),
3677                        ]
3678                        .into(),
3679                    ),
3680                    false,
3681                )),
3682                false,
3683            ),
3684            true,
3685        )]));
3686        for cdc in [false, true] {
3687            let props = WriterProperties::builder()
3688                .set_content_defined_chunking(cdc.then(Default::default))
3689                .build();
3690            for (array, expected_keys, expected_null_maps) in [
3691                (map.clone(), None, 0),
3692                (map.slice(0, 1), Some(1), 0),
3693                (map.slice(2, 1), Some(1), 0),
3694                (masked.clone(), Some(2), 1),
3695            ] {
3696                let schema = Arc::new(Schema::new(vec![Field::new(
3697                    "map",
3698                    array.data_type().clone(),
3699                    true,
3700                )]));
3701                let batch = RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap();
3702                let mut bytes = Vec::new();
3703                let mut writer =
3704                    ArrowWriter::try_new(&mut bytes, writer_schema.clone(), Some(props.clone()))
3705                        .unwrap();
3706                let result = writer.write(&batch);
3707                let Some(expected_keys) = expected_keys else {
3708                    let err = result.unwrap_err();
3709                    assert!(
3710                        err.to_string().contains("Found null")
3711                            && err.to_string().contains("required field 'key'"),
3712                        "{err}"
3713                    );
3714                    continue;
3715                };
3716                result.unwrap();
3717                writer.close().unwrap();
3718                let actual = ParquetRecordBatchReader::try_new(Bytes::from(bytes), 1024)
3719                    .unwrap()
3720                    .next()
3721                    .unwrap()
3722                    .unwrap();
3723                let actual = actual.column(0).as_map();
3724                assert_eq!(actual.null_count(), expected_null_maps);
3725                assert_eq!(
3726                    actual.keys().as_primitive::<Int32Type>().values().as_ref(),
3727                    vec![7; expected_keys]
3728                );
3729            }
3730        }
3731    }
3732
3733    #[test]
3734    fn arrow_writer_rejects_nullable_map_key_schemas_early() {
3735        let schema_with_key = |key| {
3736            Arc::new(Schema::new(vec![Field::new(
3737                "map",
3738                DataType::Map(
3739                    Arc::new(Field::new(
3740                        "entries",
3741                        DataType::Struct(
3742                            vec![key, Field::new("value", DataType::Int32, true)].into(),
3743                        ),
3744                        false,
3745                    )),
3746                    false,
3747                ),
3748                true,
3749            )]))
3750        };
3751        let required = schema_with_key(Field::new("key", DataType::Int32, false));
3752        let parquet = ArrowSchemaConverter::new().convert(&required).unwrap();
3753        let ree = DataType::RunEndEncoded(
3754            Arc::new(Field::new("run_ends", DataType::Int32, false)),
3755            Arc::new(Field::new("values", DataType::Int32, true)),
3756        );
3757        for key in [
3758            Field::new("key", DataType::Int32, true),
3759            Field::new("key", ree.clone(), false),
3760            Field::new(
3761                "key",
3762                DataType::Dictionary(Box::new(DataType::Int8), Box::new(ree)),
3763                false,
3764            ),
3765        ] {
3766            let schema = schema_with_key(key);
3767            for explicit_schema in [false, true] {
3768                for skip_metadata in [false, true] {
3769                    let mut options =
3770                        ArrowWriterOptions::new().with_skip_arrow_metadata(skip_metadata);
3771                    if explicit_schema {
3772                        options = options.with_parquet_schema(parquet.clone());
3773                    }
3774                    let mut bytes = Vec::new();
3775                    let err =
3776                        ArrowWriter::try_new_with_options(&mut bytes, schema.clone(), options)
3777                            .unwrap_err();
3778                    assert!(
3779                        err.to_string()
3780                            .contains("Map key field 'key' must be non-nullable"),
3781                        "{err}"
3782                    );
3783                    assert!(
3784                        bytes.is_empty(),
3785                        "schema rejection must precede file initialization"
3786                    );
3787                }
3788            }
3789            let file = SerializedFileWriter::new(
3790                Vec::new(),
3791                parquet.root_schema_ptr(),
3792                Arc::new(WriterProperties::default()),
3793            )
3794            .unwrap();
3795            let eager_error =
3796                ArrowRowGroupWriterFactory::try_new(&file, Arc::clone(&schema)).unwrap_err();
3797            assert!(matches!(&eager_error, ParquetError::ArrowError(_)));
3798            assert_eq!(
3799                eager_error.to_string(),
3800                "Arrow: Map key field 'key' must be non-nullable, including its encoded values"
3801            );
3802            #[expect(deprecated)] // Verify that the legacy constructor still defers errors.
3803            let factory = ArrowRowGroupWriterFactory::new(&file, schema);
3804            for row_group_index in 0..2 {
3805                let Err(err) = factory.create_column_writers(row_group_index) else {
3806                    panic!("nullable map-key schema should not create column writers")
3807                };
3808                assert!(matches!(&err, ParquetError::ArrowError(_)), "{err:?}");
3809                assert_eq!(
3810                    err.to_string(),
3811                    "Arrow: Map key field 'key' must be non-nullable, including its encoded values"
3812                );
3813            }
3814        }
3815    }
3816
3817    #[test]
3818    fn eager_row_group_validation_rejects_writer_and_chunker_count_mismatches() {
3819        let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)]));
3820        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![1]))])
3821            .unwrap();
3822        let writer = ArrowWriter::try_new(Vec::new(), schema, None).unwrap();
3823        let (_file, factory) = writer.into_serialized_writer().unwrap();
3824        let mut group = factory.create_row_group_writer(0).unwrap();
3825        assert!(
3826            group
3827                .write_with_chunkers(&batch, &mut [])
3828                .unwrap_err()
3829                .to_string()
3830                .contains("chunkers")
3831        );
3832        assert_eq!(group.buffered_rows, 0);
3833        group.writers.clear();
3834        assert!(
3835            group
3836                .write(&batch)
3837                .unwrap_err()
3838                .to_string()
3839                .contains("column writers")
3840        );
3841        assert_eq!(group.buffered_rows, 0);
3842    }
3843
3844    #[test]
3845    fn arrow_write_schema_plan_preserves_error_variants() {
3846        for error in [
3847            ParquetError::General("schema mismatch".to_owned()),
3848            ParquetError::NYI("unsupported type".to_owned()),
3849            ParquetError::ArrowError("invalid schema".to_owned()),
3850        ] {
3851            let expected_variant = std::mem::discriminant(&error);
3852            let expected_message = error.to_string();
3853            let cached = ArrowWriteSchemaPlanError::from_parquet_error(error);
3854            for _ in 0..2 {
3855                let error = cached.to_parquet_error();
3856                assert_eq!(std::mem::discriminant(&error), expected_variant);
3857                assert_eq!(error.to_string(), expected_message);
3858            }
3859        }
3860    }
3861
3862    #[test]
3863    fn arrow_write_schema_plan_caches_nested_leaf_ranges() {
3864        let nested = ArrowDataType::Struct(Fields::from(vec![
3865            Field::new("i", ArrowDataType::Int32, false),
3866            Field::new("s", ArrowDataType::Utf8, true),
3867        ]));
3868        let list = ArrowDataType::List(Arc::new(Field::new("item", ArrowDataType::Boolean, true)));
3869        let ree = ArrowDataType::RunEndEncoded(
3870            Arc::new(Field::new("run_ends", ArrowDataType::Int32, false)),
3871            Arc::new(Field::new("values", ArrowDataType::Int64, false)),
3872        );
3873        let dictionary_fsb = ArrowDataType::Dictionary(
3874            Box::new(ArrowDataType::Int8),
3875            Box::new(ArrowDataType::FixedSizeBinary(2)),
3876        );
3877        let schema = Arc::new(Schema::new(vec![
3878            Field::new("nested", nested, true),
3879            Field::new("list", list, true),
3880            Field::new("ree", ree, false),
3881            Field::new("dictionary_fsb", dictionary_fsb, true),
3882        ]));
3883        let parquet = ArrowSchemaConverter::new().convert(&schema).unwrap();
3884
3885        let plan = ArrowWriteSchemaPlan::try_new(&parquet, &schema).unwrap();
3886        let ranges = plan
3887            .fields
3888            .iter()
3889            .map(|field| field.leaf_range.clone())
3890            .collect::<Vec<_>>();
3891        assert_eq!(ranges, vec![0..2, 2..3, 3..4, 4..5]);
3892        assert_eq!(plan.leaves.len(), parquet.num_columns());
3893        assert_eq!(
3894            plan.leaves[4].physical_type(),
3895            crate::basic::Type::FIXED_LEN_BYTE_ARRAY
3896        );
3897    }
3898
3899    #[test]
3900    fn cached_schema_plan_preserves_compatible_physical_layout_alternation() {
3901        let writer_schema = Arc::new(Schema::new(vec![
3902            Field::new("number", ArrowDataType::Int32, false),
3903            Field::new("text", ArrowDataType::Utf8, false),
3904            Field::new("bytes", ArrowDataType::Binary, false),
3905        ]));
3906
3907        let dense = RecordBatch::try_new(
3908            Arc::clone(&writer_schema),
3909            vec![
3910                Arc::new(Int32Array::from(vec![1])),
3911                Arc::new(StringArray::from(vec!["a"])),
3912                Arc::new(BinaryArray::from_iter_values([b"a".as_slice()])),
3913            ],
3914        )
3915        .unwrap();
3916
3917        let number_dictionary = DictionaryArray::<Int8Type>::try_new(
3918            Int8Array::from(vec![0]),
3919            Arc::new(Int32Array::from(vec![2])),
3920        )
3921        .unwrap();
3922        let text_dictionary = DictionaryArray::<Int8Type>::try_new(
3923            Int8Array::from(vec![0]),
3924            Arc::new(StringArray::from(vec!["b"])),
3925        )
3926        .unwrap();
3927        let bytes_dictionary = DictionaryArray::<Int8Type>::try_new(
3928            Int8Array::from(vec![0]),
3929            Arc::new(BinaryArray::from_iter_values([b"b".as_slice()])),
3930        )
3931        .unwrap();
3932        let dictionary = RecordBatch::try_from_iter(vec![
3933            ("number", Arc::new(number_dictionary) as ArrayRef),
3934            ("text", Arc::new(text_dictionary) as ArrayRef),
3935            ("bytes", Arc::new(bytes_dictionary) as ArrayRef),
3936        ])
3937        .unwrap();
3938
3939        let run_ends = Int32Array::from(vec![1]);
3940        let number_ree: ArrayRef =
3941            Arc::new(Int32RunArray::try_new(&run_ends, &Int32Array::from(vec![3])).unwrap());
3942        let text_ree: ArrayRef =
3943            Arc::new(Int32RunArray::try_new(&run_ends, &StringArray::from(vec!["c"])).unwrap());
3944        let bytes_ree: ArrayRef = Arc::new(
3945            Int32RunArray::try_new(&run_ends, &BinaryArray::from_iter_values([b"c".as_slice()]))
3946                .unwrap(),
3947        );
3948        let ree = RecordBatch::try_from_iter(vec![
3949            ("number", number_ree),
3950            ("text", text_ree),
3951            ("bytes", bytes_ree),
3952        ])
3953        .unwrap();
3954
3955        let alternate = RecordBatch::try_from_iter(vec![
3956            ("number", Arc::new(Int32Array::from(vec![4])) as ArrayRef),
3957            (
3958                "text",
3959                Arc::new(LargeStringArray::from(vec!["d"])) as ArrayRef,
3960            ),
3961            (
3962                "bytes",
3963                Arc::new(BinaryViewArray::from_iter_values([b"d".as_slice()])) as ArrayRef,
3964            ),
3965        ])
3966        .unwrap();
3967
3968        let props = WriterProperties::builder()
3969            .set_max_row_group_row_count(Some(1))
3970            .build();
3971        let mut out = Vec::new();
3972        let mut writer =
3973            ArrowWriter::try_new(&mut out, Arc::clone(&writer_schema), Some(props)).unwrap();
3974        for batch in [&dense, &dictionary, &ree, &alternate] {
3975            writer.write(batch).unwrap();
3976        }
3977        writer.close().unwrap();
3978
3979        let builder = ParquetRecordBatchReaderBuilder::try_new(Bytes::from(out)).unwrap();
3980        assert_eq!(builder.metadata().num_row_groups(), 4);
3981        let mut reader = builder.build().unwrap();
3982        let actual = reader.next().unwrap().unwrap();
3983        assert_eq!(
3984            actual.column(0).as_primitive::<Int32Type>(),
3985            &Int32Array::from(vec![1, 2, 3, 4])
3986        );
3987        assert_eq!(
3988            actual.column(1).as_string::<i32>(),
3989            &StringArray::from(vec!["a", "b", "c", "d"])
3990        );
3991        assert_eq!(
3992            actual.column(2).as_binary::<i32>(),
3993            &BinaryArray::from_iter_values([
3994                b"a".as_slice(),
3995                b"b".as_slice(),
3996                b"c".as_slice(),
3997                b"d".as_slice(),
3998            ])
3999        );
4000    }
4001
4002    #[test]
4003    fn cached_schema_plan_does_not_retain_batch_arrays() {
4004        let schema = Arc::new(Schema::new(vec![Field::new(
4005            "a",
4006            ArrowDataType::Int32,
4007            false,
4008        )]));
4009        let array = Arc::new(Int32Array::from(vec![1, 2, 3]));
4010        let batch = RecordBatch::try_new(Arc::clone(&schema), vec![array.clone()]).unwrap();
4011        let mut writer = ArrowWriter::try_new(Vec::new(), schema, None).unwrap();
4012
4013        writer.write(&batch).unwrap();
4014        drop(batch);
4015        assert_eq!(
4016            Arc::strong_count(&array),
4017            1,
4018            "schema and row-group caches must not retain input ArrayRefs"
4019        );
4020        writer.close().unwrap();
4021    }
4022
4023    #[test]
4024    fn arrow_writer_page_size() {
4025        let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)]));
4026
4027        let mut builder = StringBuilder::with_capacity(100, 329 * 10_000);
4028
4029        // Generate an array of 10 unique 10 character string
4030        for i in 0..10 {
4031            let value = i
4032                .to_string()
4033                .repeat(10)
4034                .chars()
4035                .take(10)
4036                .collect::<String>();
4037
4038            builder.append_value(value);
4039        }
4040
4041        let array = Arc::new(builder.finish());
4042
4043        let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
4044
4045        let file = tempfile::tempfile().unwrap();
4046
4047        // Set everything very low so we fallback to PLAIN encoding after the first row
4048        let props = WriterProperties::builder()
4049            .set_data_page_size_limit(1)
4050            .set_dictionary_page_size_limit(1)
4051            .set_write_batch_size(1)
4052            .build();
4053
4054        let mut writer =
4055            ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props))
4056                .expect("Unable to write file");
4057        writer.write(&batch).unwrap();
4058        writer.close().unwrap();
4059
4060        let options = ReadOptionsBuilder::new().with_page_index().build();
4061        let reader =
4062            SerializedFileReader::new_with_options(file.try_clone().unwrap(), options).unwrap();
4063
4064        let column = reader.metadata().row_group(0).columns();
4065
4066        assert_eq!(column.len(), 1);
4067
4068        // We should write one row before falling back to PLAIN encoding so there should still be a
4069        // dictionary page.
4070        assert!(
4071            column[0].dictionary_page_offset().is_some(),
4072            "Expected a dictionary page"
4073        );
4074
4075        let page_index = reader
4076            .metadata()
4077            .page_index()
4078            .expect("page index should be present");
4079        let page_locations = page_index
4080            .page_locations(0, 0)
4081            .expect("page locations should exist");
4082
4083        // We should fallback to PLAIN encoding after the first row and our max page size is 1 bytes
4084        // so we expect one dictionary encoded page and then a page per row thereafter.
4085        assert_eq!(
4086            page_locations.len(),
4087            10,
4088            "Expected 10 pages but got {page_locations:#?}"
4089        );
4090    }
4091
4092    const MEDIUM_SIZE: usize = 63;
4093
4094    fn check_bloom_filter<T: AsBytes>(
4095        files: Vec<Bytes>,
4096        file_column: String,
4097        positive_values: Vec<T>,
4098        negative_values: Vec<T>,
4099    ) {
4100        files.into_iter().take(1).for_each(|file| {
4101            let file_reader = SerializedFileReader::new_with_options(
4102                file,
4103                ReadOptionsBuilder::new()
4104                    .with_reader_properties(
4105                        ReaderProperties::builder()
4106                            .set_read_bloom_filter(true)
4107                            .build(),
4108                    )
4109                    .build(),
4110            )
4111            .expect("Unable to open file as Parquet");
4112            let metadata = file_reader.metadata();
4113
4114            // Gets bloom filters from all row groups.
4115            let mut bloom_filters: Vec<_> = vec![];
4116            for (ri, row_group) in metadata.row_groups().iter().enumerate() {
4117                if let Some((column_index, _)) = row_group
4118                    .columns()
4119                    .iter()
4120                    .enumerate()
4121                    .find(|(_, column)| column.column_path().string() == file_column)
4122                {
4123                    let row_group_reader = file_reader
4124                        .get_row_group(ri)
4125                        .expect("Unable to read row group");
4126                    if let Some(sbbf) = row_group_reader.get_column_bloom_filter(column_index) {
4127                        bloom_filters.push(sbbf.clone());
4128                    } else {
4129                        panic!("No bloom filter for column named {file_column} found");
4130                    }
4131                } else {
4132                    panic!("No column named {file_column} found");
4133                }
4134            }
4135
4136            positive_values.iter().for_each(|value| {
4137                let found = bloom_filters.iter().find(|sbbf| sbbf.check(value));
4138                assert!(
4139                    found.is_some(),
4140                    "{}",
4141                    format!("Value {:?} should be in bloom filter", value.as_bytes())
4142                );
4143            });
4144
4145            negative_values.iter().for_each(|value| {
4146                let found = bloom_filters.iter().find(|sbbf| sbbf.check(value));
4147                assert!(
4148                    found.is_none(),
4149                    "{}",
4150                    format!("Value {:?} should not be in bloom filter", value.as_bytes())
4151                );
4152            });
4153        });
4154    }
4155
4156    #[test]
4157    #[cfg_attr(miri, ignore)] // Takes too long
4158    fn bool_large_single_column() {
4159        let values = Arc::new(
4160            [None, Some(true), Some(false)]
4161                .iter()
4162                .cycle()
4163                .copied()
4164                .take(200_000)
4165                .collect::<BooleanArray>(),
4166        );
4167        let schema = Schema::new(vec![Field::new("col", values.data_type().clone(), true)]);
4168        let expected_batch = RecordBatch::try_new(Arc::new(schema), vec![values]).unwrap();
4169        let file = tempfile::tempfile().unwrap();
4170
4171        let mut writer =
4172            ArrowWriter::try_new(file.try_clone().unwrap(), expected_batch.schema(), None)
4173                .expect("Unable to write file");
4174        writer.write(&expected_batch).unwrap();
4175        writer.close().unwrap();
4176    }
4177
4178    #[test]
4179    fn check_page_offset_index_with_nan() {
4180        let values = Arc::new(Float64Array::from(vec![f64::NAN; 10]));
4181        let schema = Schema::new(vec![Field::new("col", DataType::Float64, true)]);
4182        let batch = RecordBatch::try_new(Arc::new(schema), vec![values]).unwrap();
4183
4184        let mut out = Vec::with_capacity(1024);
4185        let mut writer =
4186            ArrowWriter::try_new(&mut out, batch.schema(), None).expect("Unable to write file");
4187        writer.write(&batch).unwrap();
4188        let file_meta_data = writer.close().unwrap();
4189        for row_group in file_meta_data.row_groups() {
4190            for column in row_group.columns() {
4191                assert!(column.offset_index_offset().is_some());
4192                assert!(column.offset_index_length().is_some());
4193                assert!(column.column_index_offset().is_some());
4194                assert!(column.column_index_length().is_some());
4195            }
4196        }
4197        if let Some(page_index) = file_meta_data.page_index() {
4198            for rg in 0..file_meta_data.num_row_groups() {
4199                for col in 0..file_meta_data.row_group(rg).num_columns() {
4200                    let idx = page_index
4201                        .column_index(rg, col)
4202                        .expect("column index should exist");
4203                    assert!(idx.nan_counts().is_some());
4204                    let ColumnIndexMetaData::DOUBLE(float_idx) = idx else {
4205                        panic!("expected double statistics")
4206                    };
4207                    for i in 0..idx.num_pages() as usize {
4208                        assert_eq!(float_idx.nan_count(i), Some(10));
4209                        assert_eq!(
4210                            f64::NAN.total_cmp(float_idx.min_value(i).unwrap()),
4211                            Ordering::Equal
4212                        );
4213                        assert_eq!(
4214                            f64::NAN.total_cmp(float_idx.max_value(i).unwrap()),
4215                            Ordering::Equal
4216                        );
4217                    }
4218                }
4219            }
4220        } else {
4221            panic!("page index should be present");
4222        }
4223    }
4224
4225    #[test]
4226    fn check_page_offset_index_with_mixed_nan() {
4227        let schema = Arc::new(Schema::new(vec![Field::new(
4228            "col",
4229            DataType::Float64,
4230            true,
4231        )]));
4232
4233        let mut out = Vec::with_capacity(1024);
4234        let props = WriterProperties::builder()
4235            .set_data_page_row_count_limit(10)
4236            .build();
4237        let mut writer = ArrowWriter::try_new(&mut out, schema.clone(), Some(props))
4238            .expect("Unable to write file");
4239
4240        // write a page of all NaN (since batch min and max are NaN, global min/max are NaN)
4241        let values = Arc::new(Float64Array::from(vec![f64::NAN; 10]));
4242        let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
4243        writer.write(&batch).unwrap();
4244
4245        // write a page of all -NaN (batch min/max is -NaN, should update global min to -NaN)
4246        let values = Arc::new(Float64Array::from(vec![-f64::NAN; 10]));
4247        let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
4248        writer.write(&batch).unwrap();
4249
4250        // write a page of all 0 (non-NaN should override global min/max, now 0/0)
4251        let values = Arc::new(Float64Array::from(vec![0_f64; 10]));
4252        let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
4253        writer.write(&batch).unwrap();
4254
4255        // write a mixed page (should now have min -1, max 1)
4256        let values = Arc::new(Float64Array::from(vec![
4257            -1.0,
4258            0.0,
4259            f64::NAN,
4260            -f64::NAN,
4261            1.0,
4262        ]));
4263        let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
4264        writer.write(&batch).unwrap();
4265
4266        let file_meta_data = writer.close().unwrap();
4267
4268        // check the column chunk stats are correct
4269        let col_stats = file_meta_data
4270            .row_group(0)
4271            .column(0)
4272            .statistics()
4273            .expect("missing column chunk statistics");
4274
4275        assert_eq!(col_stats.nan_count_opt(), Some(22));
4276        assert_eq!(col_stats.min_bytes_opt(), Some((-1.0f64).as_bytes()));
4277        assert_eq!(col_stats.max_bytes_opt(), Some(1.0f64.as_bytes()));
4278
4279        assert!(file_meta_data.page_index().is_some());
4280        let col_idx = &file_meta_data.page_index().unwrap().column_index(0, 0);
4281        assert_eq!(col_idx.as_ref().unwrap().num_pages(), 4);
4282
4283        // test each page
4284        let Some(ColumnIndexMetaData::DOUBLE(float_idx)) = col_idx else {
4285            panic!("expected double statistics")
4286        };
4287
4288        assert_eq!(float_idx.nan_counts, Some(vec![10, 10, 0, 2]));
4289        assert_eq!(
4290            f64::NAN.total_cmp(float_idx.min_value(0).unwrap()),
4291            Ordering::Equal
4292        );
4293        assert_eq!(
4294            f64::NAN.total_cmp(float_idx.max_value(0).unwrap()),
4295            Ordering::Equal
4296        );
4297        assert_eq!(
4298            (-f64::NAN).total_cmp(float_idx.min_value(1).unwrap()),
4299            Ordering::Equal
4300        );
4301        assert_eq!(
4302            (-f64::NAN).total_cmp(float_idx.max_value(1).unwrap()),
4303            Ordering::Equal
4304        );
4305        assert_eq!(float_idx.min_value(2), Some(&0.0));
4306        assert_eq!(float_idx.max_value(2), Some(&0.0));
4307        assert_eq!(float_idx.min_value(3), Some(&-1.0));
4308        assert_eq!(float_idx.max_value(3), Some(&1.0));
4309    }
4310
4311    #[test]
4312    #[should_panic(
4313        expected = "Attempting to write an Arrow interval type MonthDayNano to parquet that is not yet implemented"
4314    )]
4315    fn interval_month_day_nano_single_column() {
4316        required_and_optional::<IntervalMonthDayNanoArray, _>(vec![
4317            IntervalMonthDayNano::new(0, 1, 5),
4318            IntervalMonthDayNano::new(0, 3, 2),
4319            IntervalMonthDayNano::new(3, -2, -5),
4320            IntervalMonthDayNano::new(-200, 4, -1),
4321        ]);
4322    }
4323
4324    #[test]
4325    #[cfg_attr(miri, ignore)] // Takes too long
4326    fn i32_column_bloom_filter_at_end() {
4327        let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
4328        let files = RoundTripTest::new(array)
4329            .with_nullable(false)
4330            .with_bloom_filter(true)
4331            .with_bloom_filter_position(BloomFilterPosition::End)
4332            .run();
4333
4334        check_bloom_filter(
4335            files,
4336            "col".to_string(),
4337            (0..SMALL_SIZE as i32).collect(),
4338            (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
4339        );
4340    }
4341
4342    #[test]
4343    #[cfg_attr(miri, ignore)] // Takes too long
4344    fn i32_column_bloom_filter() {
4345        let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
4346        let files = RoundTripTest::new(array)
4347            .with_nullable(false)
4348            .with_bloom_filter(true)
4349            .run();
4350
4351        check_bloom_filter(
4352            files,
4353            "col".to_string(),
4354            (0..SMALL_SIZE as i32).collect(),
4355            (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
4356        );
4357    }
4358
4359    fn write_with_bloom_filter(array: ArrayRef, dictionary_page_size_limit: usize) -> Bytes {
4360        let schema = Arc::new(Schema::new(vec![Field::new(
4361            "col",
4362            array.data_type().clone(),
4363            false,
4364        )]));
4365        let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap();
4366        let props = WriterProperties::builder()
4367            .set_dictionary_enabled(true)
4368            .set_dictionary_page_size_limit(dictionary_page_size_limit)
4369            .set_write_batch_size(256)
4370            .set_bloom_filter_enabled(true)
4371            .build();
4372        let mut buf = Vec::new();
4373        let mut writer = ArrowWriter::try_new(&mut buf, schema, Some(props)).unwrap();
4374        writer.write(&batch).unwrap();
4375        writer.close().unwrap();
4376        Bytes::from(buf)
4377    }
4378
4379    fn data_page_encoding_mask(file: &Bytes) -> EncodingMask {
4380        let metadata = ParquetMetaDataReader::new().parse_and_finish(file).unwrap();
4381        *metadata
4382            .row_group(0)
4383            .column(0)
4384            .page_encoding_stats_mask()
4385            .unwrap()
4386    }
4387
4388    /// While a column is dictionary encoded the bloom filter is populated from the dictionary
4389    /// when it is flushed, so a chunk that stays dictionary encoded must still contain every value.
4390    #[test]
4391    fn string_column_bloom_filter_populated_from_dictionary() {
4392        let values: Vec<String> = (0..2000).map(|i| format!("value-{}", i % 10)).collect();
4393        let array = Arc::new(StringArray::from_iter_values(&values));
4394        let file = write_with_bloom_filter(array, 1024 * 1024);
4395        assert!(data_page_encoding_mask(&file).is_only(Encoding::RLE_DICTIONARY));
4396
4397        check_bloom_filter(
4398            vec![file],
4399            "col".to_string(),
4400            (0..10).map(|i| format!("value-{i}").into_bytes()).collect(),
4401            (10..20)
4402                .map(|i| format!("value-{i}").into_bytes())
4403                .collect(),
4404        );
4405    }
4406
4407    /// After falling back from dictionary encoding the filter holds the dictionary's values
4408    /// and every value written plain afterwards.
4409    #[test]
4410    fn string_column_bloom_filter_across_dictionary_fallback() {
4411        let values: Vec<String> = (0..2000).map(|i| format!("value-{i}")).collect();
4412        let array = Arc::new(StringArray::from_iter_values(&values));
4413        let file = write_with_bloom_filter(array, 1024);
4414        let encodings = data_page_encoding_mask(&file);
4415        assert!(
4416            encodings.is_set(Encoding::RLE_DICTIONARY) && encodings.is_set(Encoding::PLAIN),
4417            "expected dictionary and plain data pages, got {encodings:?}"
4418        );
4419
4420        check_bloom_filter(
4421            vec![file],
4422            "col".to_string(),
4423            values.into_iter().map(String::into_bytes).collect(),
4424            (2000..2010)
4425                .map(|i| format!("value-{i}").into_bytes())
4426                .collect(),
4427        );
4428    }
4429
4430    #[test]
4431    fn i64_column_bloom_filter_populated_from_dictionary() {
4432        let array = Arc::new(Int64Array::from_iter_values((0..2000).map(|i| i % 10)));
4433        let file = write_with_bloom_filter(array, 1024 * 1024);
4434        assert!(data_page_encoding_mask(&file).is_only(Encoding::RLE_DICTIONARY));
4435
4436        check_bloom_filter(
4437            vec![file],
4438            "col".to_string(),
4439            (0..10i64).collect(),
4440            (10..20i64).collect(),
4441        );
4442    }
4443
4444    #[test]
4445    fn i64_column_bloom_filter_across_dictionary_fallback() {
4446        let array = Arc::new(Int64Array::from_iter_values(0..2000i64));
4447        let file = write_with_bloom_filter(array, 1024);
4448        let encodings = data_page_encoding_mask(&file);
4449        assert!(
4450            encodings.is_set(Encoding::RLE_DICTIONARY) && encodings.is_set(Encoding::PLAIN),
4451            "expected dictionary and plain data pages, got {encodings:?}"
4452        );
4453
4454        check_bloom_filter(
4455            vec![file],
4456            "col".to_string(),
4457            (0..2000i64).collect(),
4458            (2000..2010i64).collect(),
4459        );
4460    }
4461
4462    /// Test that bloom filter folding produces correct results even when
4463    /// the configured NDV differs significantly from actual NDV.
4464    /// A large NDV means a larger initial filter that gets folded down;
4465    /// a small NDV means a smaller initial filter.
4466    #[test]
4467    #[cfg_attr(miri, ignore)] // Takes too long
4468    fn i32_column_bloom_filter_fixed_ndv() {
4469        let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
4470
4471        // NDV much larger than actual distinct values — tests folding a large filter down
4472        let files = RoundTripTest::new(array.clone())
4473            .with_nullable(false)
4474            .with_bloom_filter(true)
4475            .with_bloom_filter_ndv(1_000_000)
4476            .run();
4477
4478        check_bloom_filter(
4479            files,
4480            "col".to_string(),
4481            (0..SMALL_SIZE as i32).collect(),
4482            (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
4483        );
4484
4485        // NDV smaller than actual distinct values — tests the underestimate path
4486        let files = RoundTripTest::new(array)
4487            .with_nullable(false)
4488            .with_bloom_filter(true)
4489            .with_bloom_filter_ndv(3)
4490            .run();
4491
4492        check_bloom_filter(
4493            files,
4494            "col".to_string(),
4495            (0..SMALL_SIZE as i32).collect(),
4496            (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
4497        );
4498    }
4499
4500    #[test]
4501    #[cfg_attr(miri, ignore)] // Takes too long
4502    fn binary_column_bloom_filter() {
4503        let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
4504        let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
4505        let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
4506
4507        let array = Arc::new(BinaryArray::from_iter_values(many_vecs_iter));
4508        let files = RoundTripTest::new(array)
4509            .with_nullable(false)
4510            .with_bloom_filter(true)
4511            .run();
4512
4513        check_bloom_filter(
4514            files,
4515            "col".to_string(),
4516            many_vecs,
4517            vec![vec![(SMALL_SIZE + 1) as u8]],
4518        );
4519    }
4520
4521    #[test]
4522    #[cfg_attr(miri, ignore)] // Takes too long
4523    fn empty_string_null_column_bloom_filter() {
4524        let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
4525        let raw_strs = raw_values.iter().map(|s| s.as_str());
4526
4527        let array = Arc::new(StringArray::from_iter_values(raw_strs));
4528        let files = RoundTripTest::new(array)
4529            .with_nullable(false)
4530            .with_bloom_filter(true)
4531            .run();
4532
4533        let optional_raw_values: Vec<_> = raw_values
4534            .iter()
4535            .enumerate()
4536            .filter_map(|(i, v)| if i % 2 == 0 { None } else { Some(v.as_str()) })
4537            .collect();
4538        // For null slots, empty string should not be in bloom filter.
4539        check_bloom_filter(files, "col".to_string(), optional_raw_values, vec![""]);
4540    }
4541
4542    #[test]
4543    fn list_and_map_coerced_names() {
4544        // Create map and list with non-Parquet naming
4545        let list_field =
4546            Field::new_list("my_list", Field::new("item", DataType::Int32, false), false);
4547        let map_field = Field::new_map(
4548            "my_map",
4549            "my_entries",
4550            Field::new("my_keys", DataType::Int32, false),
4551            Field::new("my_values", DataType::Int32, true),
4552            false,
4553            true,
4554        );
4555
4556        let list_array = create_random_array(&list_field, 100, 0.0, 0.0).unwrap();
4557        let map_array = create_random_array(&map_field, 100, 0.0, 0.0).unwrap();
4558
4559        let arrow_schema = Arc::new(Schema::new(vec![list_field, map_field]));
4560
4561        // Write data to Parquet but coerce names to match spec
4562        let props = Some(WriterProperties::builder().set_coerce_types(true).build());
4563        let file = tempfile::tempfile().unwrap();
4564        let mut writer =
4565            ArrowWriter::try_new(file.try_clone().unwrap(), arrow_schema.clone(), props).unwrap();
4566
4567        let batch = RecordBatch::try_new(arrow_schema, vec![list_array, map_array]).unwrap();
4568        writer.write(&batch).unwrap();
4569        let file_metadata = writer.close().unwrap();
4570
4571        let schema = file_metadata.file_metadata().schema();
4572        // Coerced name of "item" should be "element"
4573        let list_field = &schema.get_fields()[0].get_fields()[0];
4574        assert_eq!(list_field.get_fields()[0].name(), "element");
4575
4576        let map_field = &schema.get_fields()[1].get_fields()[0];
4577        // Coerced name of "entries" should be "key_value"
4578        assert_eq!(map_field.name(), "key_value");
4579        // Coerced name of "my_keys" should be "key"
4580        assert_eq!(map_field.get_fields()[0].name(), "key");
4581        // Coerced name of "my_values" should be "value"
4582        assert_eq!(map_field.get_fields()[1].name(), "value");
4583
4584        // Double check schema after reading from the file
4585        let reader = SerializedFileReader::new(file).unwrap();
4586        let file_schema = reader.metadata().file_metadata().schema();
4587        let fields = file_schema.get_fields();
4588        let list_field = &fields[0].get_fields()[0];
4589        assert_eq!(list_field.get_fields()[0].name(), "element");
4590        let map_field = &fields[1].get_fields()[0];
4591        assert_eq!(map_field.name(), "key_value");
4592        assert_eq!(map_field.get_fields()[0].name(), "key");
4593        assert_eq!(map_field.get_fields()[1].name(), "value");
4594    }
4595
4596    #[test]
4597    #[cfg_attr(miri, ignore)] // Takes too long
4598    fn fallback_flush_data_page() {
4599        //tests if the Fallback::flush_data_page clears all buffers correctly
4600        let raw_values: Vec<_> = (0..MEDIUM_SIZE).map(|i| i.to_string()).collect();
4601        let values = Arc::new(StringArray::from(raw_values));
4602        let encodings = vec![
4603            Encoding::DELTA_BYTE_ARRAY,
4604            Encoding::DELTA_LENGTH_BYTE_ARRAY,
4605        ];
4606        let data_type = values.data_type().clone();
4607        let schema = Arc::new(Schema::new(vec![Field::new("col", data_type, false)]));
4608        let expected_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4609
4610        let row_group_sizes = [1024, SMALL_SIZE, SMALL_SIZE / 2, SMALL_SIZE / 2 + 1, 10];
4611        let data_page_size_limit: usize = 32;
4612        let write_batch_size: usize = 16;
4613
4614        for encoding in &encodings {
4615            for row_group_size in row_group_sizes {
4616                let props = WriterProperties::builder()
4617                    .set_writer_version(WriterVersion::PARQUET_2_0)
4618                    .set_max_row_group_row_count(Some(row_group_size))
4619                    .set_dictionary_enabled(false)
4620                    .set_encoding(*encoding)
4621                    .set_data_page_size_limit(data_page_size_limit)
4622                    .set_write_batch_size(write_batch_size)
4623                    .build();
4624
4625                roundtrip_opts_with_array_validation(&expected_batch, props, |a, b| {
4626                    let string_array_a = StringArray::from(a.clone());
4627                    let string_array_b = StringArray::from(b.clone());
4628                    let vec_a: Vec<&str> = string_array_a.iter().map(|v| v.unwrap()).collect();
4629                    let vec_b: Vec<&str> = string_array_b.iter().map(|v| v.unwrap()).collect();
4630                    assert_eq!(
4631                        vec_a, vec_b,
4632                        "failed for encoder: {encoding:?} and row_group_size: {row_group_size:?}"
4633                    );
4634                });
4635            }
4636        }
4637    }
4638
4639    #[test]
4640    #[cfg_attr(miri, ignore)] // Takes too long
4641    fn u32_min_max() {
4642        // check values roundtrip through parquet
4643        let src = [
4644            u32::MIN,
4645            1,
4646            (i32::MAX as u32) - 1,
4647            i32::MAX as u32,
4648            (i32::MAX as u32) + 1,
4649            u32::MAX - 1,
4650            u32::MAX,
4651        ];
4652        let values = Arc::new(UInt32Array::from_iter_values(src.iter().copied()));
4653        let files = RoundTripTest::new(values).with_nullable(false).run();
4654
4655        for file in files {
4656            // check statistics are valid
4657            let reader = SerializedFileReader::new(file).unwrap();
4658            let metadata = reader.metadata();
4659
4660            let mut row_offset = 0;
4661            for row_group in metadata.row_groups() {
4662                assert_eq!(row_group.num_columns(), 1);
4663                let column = row_group.column(0);
4664
4665                let num_values = column.num_values() as usize;
4666                let src_slice = &src[row_offset..row_offset + num_values];
4667                row_offset += column.num_values() as usize;
4668
4669                let stats = column.statistics().unwrap();
4670                if let Statistics::Int32(stats) = stats {
4671                    assert_eq!(
4672                        *stats.min_opt().unwrap() as u32,
4673                        *src_slice.iter().min().unwrap()
4674                    );
4675                    assert_eq!(
4676                        *stats.max_opt().unwrap() as u32,
4677                        *src_slice.iter().max().unwrap()
4678                    );
4679                } else {
4680                    panic!("Statistics::Int32 missing")
4681                }
4682            }
4683        }
4684    }
4685
4686    #[test]
4687    #[cfg_attr(miri, ignore)] // Takes too long
4688    fn u64_min_max() {
4689        // check values roundtrip through parquet
4690        let src = [
4691            u64::MIN,
4692            1,
4693            (i64::MAX as u64) - 1,
4694            i64::MAX as u64,
4695            (i64::MAX as u64) + 1,
4696            u64::MAX - 1,
4697            u64::MAX,
4698        ];
4699        let values = Arc::new(UInt64Array::from_iter_values(src.iter().copied()));
4700        let files = RoundTripTest::new(values).with_nullable(false).run();
4701
4702        for file in files {
4703            // check statistics are valid
4704            let reader = SerializedFileReader::new(file).unwrap();
4705            let metadata = reader.metadata();
4706
4707            let mut row_offset = 0;
4708            for row_group in metadata.row_groups() {
4709                assert_eq!(row_group.num_columns(), 1);
4710                let column = row_group.column(0);
4711
4712                let num_values = column.num_values() as usize;
4713                let src_slice = &src[row_offset..row_offset + num_values];
4714                row_offset += column.num_values() as usize;
4715
4716                let stats = column.statistics().unwrap();
4717                if let Statistics::Int64(stats) = stats {
4718                    assert_eq!(
4719                        *stats.min_opt().unwrap() as u64,
4720                        *src_slice.iter().min().unwrap()
4721                    );
4722                    assert_eq!(
4723                        *stats.max_opt().unwrap() as u64,
4724                        *src_slice.iter().max().unwrap()
4725                    );
4726                } else {
4727                    panic!("Statistics::Int64 missing")
4728                }
4729            }
4730        }
4731    }
4732
4733    #[test]
4734    #[cfg_attr(miri, ignore)] // Takes too long
4735    fn statistics_null_counts_only_nulls() {
4736        // check that null-count statistics for "only NULL"-columns are correct
4737        let values = Arc::new(UInt64Array::from(vec![None, None]));
4738        let files = RoundTripTest::new(values).run();
4739
4740        for file in files {
4741            // check statistics are valid
4742            let reader = SerializedFileReader::new(file).unwrap();
4743            let metadata = reader.metadata();
4744            assert_eq!(metadata.num_row_groups(), 1);
4745            let row_group = metadata.row_group(0);
4746            assert_eq!(row_group.num_columns(), 1);
4747            let column = row_group.column(0);
4748            let stats = column.statistics().unwrap();
4749            assert_eq!(stats.null_count_opt(), Some(2));
4750        }
4751    }
4752
4753    fn row_group_sizes(metadata: &ParquetMetaData) -> Vec<i64> {
4754        metadata.row_groups().iter().map(|x| x.num_rows()).collect()
4755    }
4756
4757    #[test]
4758    fn test_aggregates_records() {
4759        let arrays = [
4760            Int32Array::from((0..100).collect::<Vec<_>>()),
4761            Int32Array::from((0..50).collect::<Vec<_>>()),
4762            Int32Array::from((200..500).collect::<Vec<_>>()),
4763        ];
4764
4765        let schema = Arc::new(Schema::new(vec![Field::new(
4766            "int",
4767            ArrowDataType::Int32,
4768            false,
4769        )]));
4770
4771        let file = tempfile::tempfile().unwrap();
4772
4773        let props = WriterProperties::builder()
4774            .set_max_row_group_row_count(Some(200))
4775            .build();
4776
4777        let mut writer =
4778            ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
4779
4780        for array in arrays {
4781            let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
4782            writer.write(&batch).unwrap();
4783        }
4784
4785        writer.close().unwrap();
4786
4787        let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
4788        assert_eq!(&row_group_sizes(builder.metadata()), &[200, 200, 50]);
4789
4790        let batches = builder
4791            .with_batch_size(100)
4792            .build()
4793            .unwrap()
4794            .collect::<ArrowResult<Vec<_>>>()
4795            .unwrap();
4796
4797        assert_eq!(batches.len(), 5);
4798        assert!(batches.iter().all(|x| x.num_columns() == 1));
4799
4800        let batch_sizes: Vec<_> = batches.iter().map(|x| x.num_rows()).collect();
4801
4802        assert_eq!(&batch_sizes, &[100, 100, 100, 100, 50]);
4803
4804        let values: Vec<_> = batches
4805            .iter()
4806            .flat_map(|x| {
4807                x.column(0)
4808                    .as_any()
4809                    .downcast_ref::<Int32Array>()
4810                    .unwrap()
4811                    .values()
4812                    .iter()
4813                    .copied()
4814            })
4815            .collect();
4816
4817        let expected_values: Vec<_> = [0..100, 0..50, 200..500].into_iter().flatten().collect();
4818        assert_eq!(&values, &expected_values)
4819    }
4820
4821    #[test]
4822    fn complex_aggregate() {
4823        // Tests aggregating nested data
4824        let field_a = Arc::new(Field::new("leaf_a", DataType::Int32, false));
4825        let field_b = Arc::new(Field::new("leaf_b", DataType::Int32, true));
4826        let struct_a = Arc::new(Field::new(
4827            "struct_a",
4828            DataType::Struct(vec![field_a.clone(), field_b.clone()].into()),
4829            true,
4830        ));
4831
4832        let list_a = Arc::new(Field::new("list", DataType::List(struct_a), true));
4833        let struct_b = Arc::new(Field::new(
4834            "struct_b",
4835            DataType::Struct(vec![list_a.clone()].into()),
4836            false,
4837        ));
4838
4839        let schema = Arc::new(Schema::new(vec![struct_b]));
4840
4841        // create nested data
4842        let field_a_array = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
4843        let field_b_array =
4844            Int32Array::from_iter(vec![Some(1), None, Some(2), None, None, Some(6)]);
4845
4846        let struct_a_array = StructArray::from(vec![
4847            (field_a.clone(), Arc::new(field_a_array) as ArrayRef),
4848            (field_b.clone(), Arc::new(field_b_array) as ArrayRef),
4849        ]);
4850
4851        let list_data = ArrayDataBuilder::new(list_a.data_type().clone())
4852            .len(5)
4853            .add_buffer(Buffer::from_iter(vec![
4854                0_i32, 1_i32, 1_i32, 3_i32, 3_i32, 5_i32,
4855            ]))
4856            .null_bit_buffer(Some(Buffer::from_iter(vec![
4857                true, false, true, false, true,
4858            ])))
4859            .child_data(vec![struct_a_array.into_data()])
4860            .build()
4861            .unwrap();
4862
4863        let list_a_array = Arc::new(ListArray::from(list_data)) as ArrayRef;
4864        let struct_b_array = StructArray::from(vec![(list_a.clone(), list_a_array)]);
4865
4866        let batch1 =
4867            RecordBatch::try_from_iter(vec![("struct_b", Arc::new(struct_b_array) as ArrayRef)])
4868                .unwrap();
4869
4870        let field_a_array = Int32Array::from(vec![6, 7, 8, 9, 10]);
4871        let field_b_array = Int32Array::from_iter(vec![None, None, None, Some(1), None]);
4872
4873        let struct_a_array = StructArray::from(vec![
4874            (field_a, Arc::new(field_a_array) as ArrayRef),
4875            (field_b, Arc::new(field_b_array) as ArrayRef),
4876        ]);
4877
4878        let list_data = ArrayDataBuilder::new(list_a.data_type().clone())
4879            .len(2)
4880            .add_buffer(Buffer::from_iter(vec![0_i32, 4_i32, 5_i32]))
4881            .child_data(vec![struct_a_array.into_data()])
4882            .build()
4883            .unwrap();
4884
4885        let list_a_array = Arc::new(ListArray::from(list_data)) as ArrayRef;
4886        let struct_b_array = StructArray::from(vec![(list_a, list_a_array)]);
4887
4888        let batch2 =
4889            RecordBatch::try_from_iter(vec![("struct_b", Arc::new(struct_b_array) as ArrayRef)])
4890                .unwrap();
4891
4892        let batches = &[batch1, batch2];
4893
4894        // Verify data is as expected
4895
4896        let expected = r"
4897            +-------------------------------------------------------------------------------------------------------+
4898            | struct_b                                                                                              |
4899            +-------------------------------------------------------------------------------------------------------+
4900            | {list: [{leaf_a: 1, leaf_b: 1}]}                                                                      |
4901            | {list: }                                                                                              |
4902            | {list: [{leaf_a: 2, leaf_b: }, {leaf_a: 3, leaf_b: 2}]}                                               |
4903            | {list: }                                                                                              |
4904            | {list: [{leaf_a: 4, leaf_b: }, {leaf_a: 5, leaf_b: }]}                                                |
4905            | {list: [{leaf_a: 6, leaf_b: }, {leaf_a: 7, leaf_b: }, {leaf_a: 8, leaf_b: }, {leaf_a: 9, leaf_b: 1}]} |
4906            | {list: [{leaf_a: 10, leaf_b: }]}                                                                      |
4907            +-------------------------------------------------------------------------------------------------------+
4908        ".trim().split('\n').map(|x| x.trim()).collect::<Vec<_>>().join("\n");
4909
4910        let actual = pretty_format_batches(batches).unwrap().to_string();
4911        assert_eq!(actual, expected);
4912
4913        // Write data
4914        let file = tempfile::tempfile().unwrap();
4915        let props = WriterProperties::builder()
4916            .set_max_row_group_row_count(Some(6))
4917            .build();
4918
4919        let mut writer =
4920            ArrowWriter::try_new(file.try_clone().unwrap(), schema, Some(props)).unwrap();
4921
4922        for batch in batches {
4923            writer.write(batch).unwrap();
4924        }
4925        writer.close().unwrap();
4926
4927        // Read Data
4928        // Should have written entire first batch and first row of second to the first row group
4929        // leaving a single row in the second row group
4930
4931        let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
4932        assert_eq!(&row_group_sizes(builder.metadata()), &[6, 1]);
4933
4934        let batches = builder
4935            .with_batch_size(2)
4936            .build()
4937            .unwrap()
4938            .collect::<ArrowResult<Vec<_>>>()
4939            .unwrap();
4940
4941        assert_eq!(batches.len(), 4);
4942        let batch_counts: Vec<_> = batches.iter().map(|x| x.num_rows()).collect();
4943        assert_eq!(&batch_counts, &[2, 2, 2, 1]);
4944
4945        let actual = pretty_format_batches(&batches).unwrap().to_string();
4946        assert_eq!(actual, expected);
4947    }
4948
4949    #[test]
4950    fn test_arrow_writer_metadata() {
4951        let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
4952        let file_schema = batch_schema.clone().with_metadata([("foo", "bar")]);
4953
4954        let batch = RecordBatch::try_new(
4955            Arc::new(batch_schema),
4956            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
4957        )
4958        .unwrap();
4959
4960        let mut buf = Vec::with_capacity(1024);
4961        let mut writer = ArrowWriter::try_new(&mut buf, Arc::new(file_schema), None).unwrap();
4962        writer.write(&batch).unwrap();
4963        writer.close().unwrap();
4964    }
4965
4966    #[test]
4967    fn in_progress_accounting() {
4968        // define schema
4969        let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
4970
4971        // create some data
4972        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
4973
4974        // build a record batch
4975        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
4976
4977        let mut writer = ArrowWriter::try_new(vec![], batch.schema(), None).unwrap();
4978
4979        // starts empty
4980        assert_eq!(writer.in_progress_size(), 0);
4981        assert_eq!(writer.in_progress_rows(), 0);
4982        assert_eq!(writer.memory_size(), 0);
4983        assert_eq!(writer.bytes_written(), 4); // Initial header
4984        writer.write(&batch).unwrap();
4985
4986        // updated on write
4987        let initial_size = writer.in_progress_size();
4988        assert!(initial_size > 0);
4989        assert_eq!(writer.in_progress_rows(), 5);
4990        let initial_memory = writer.memory_size();
4991        assert!(initial_memory > 0);
4992        // memory estimate is larger than estimated encoded size
4993        assert!(
4994            initial_size <= initial_memory,
4995            "{initial_size} <= {initial_memory}"
4996        );
4997
4998        // updated on second write
4999        writer.write(&batch).unwrap();
5000        assert!(writer.in_progress_size() > initial_size);
5001        assert_eq!(writer.in_progress_rows(), 10);
5002        assert!(writer.memory_size() > initial_memory);
5003        assert!(
5004            writer.in_progress_size() <= writer.memory_size(),
5005            "in_progress_size {} <= memory_size {}",
5006            writer.in_progress_size(),
5007            writer.memory_size()
5008        );
5009
5010        // in progress tracking is cleared, but the overall data written is updated
5011        let pre_flush_bytes_written = writer.bytes_written();
5012        writer.flush().unwrap();
5013        assert_eq!(writer.in_progress_size(), 0);
5014        assert_eq!(writer.memory_size(), 0);
5015        assert!(writer.bytes_written() > pre_flush_bytes_written);
5016
5017        writer.close().unwrap();
5018    }
5019
5020    #[test]
5021    fn test_writer_all_null() {
5022        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
5023        let b = Int32Array::new(vec![0; 5].into(), Some(NullBuffer::new_null(5)));
5024        let batch = RecordBatch::try_from_iter(vec![
5025            ("a", Arc::new(a) as ArrayRef),
5026            ("b", Arc::new(b) as ArrayRef),
5027        ])
5028        .unwrap();
5029
5030        let mut buf = Vec::with_capacity(1024);
5031        let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), None).unwrap();
5032        writer.write(&batch).unwrap();
5033        writer.close().unwrap();
5034
5035        let bytes = Bytes::from(buf);
5036        let options = ReadOptionsBuilder::new().with_page_index().build();
5037        let reader = SerializedFileReader::new_with_options(bytes, options).unwrap();
5038        let index = reader.metadata().page_index().unwrap();
5039
5040        assert_eq!(index.num_data_pages(0, 0), Some(1)); // 1 page
5041        assert_eq!(index.num_data_pages(0, 1), Some(1)); // 1 page
5042    }
5043
5044    #[test]
5045    fn test_disabled_statistics_with_page() {
5046        let file_schema = Schema::new(vec![
5047            Field::new("a", DataType::Utf8, true),
5048            Field::new("b", DataType::Utf8, true),
5049        ]);
5050        let file_schema = Arc::new(file_schema);
5051
5052        let batch = RecordBatch::try_new(
5053            file_schema.clone(),
5054            vec![
5055                Arc::new(StringArray::from(vec!["a", "b", "c", "d"])) as _,
5056                Arc::new(StringArray::from(vec!["w", "x", "y", "z"])) as _,
5057            ],
5058        )
5059        .unwrap();
5060
5061        let props = WriterProperties::builder()
5062            .set_statistics_enabled(EnabledStatistics::None)
5063            .set_column_statistics_enabled("a".into(), EnabledStatistics::Page)
5064            .build();
5065
5066        let mut buf = Vec::with_capacity(1024);
5067        let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), Some(props)).unwrap();
5068        writer.write(&batch).unwrap();
5069
5070        let metadata = writer.close().unwrap();
5071        assert_eq!(metadata.num_row_groups(), 1);
5072        let row_group = metadata.row_group(0);
5073        assert_eq!(row_group.num_columns(), 2);
5074        // Column "a" has both offset and column index, as requested
5075        assert!(row_group.column(0).offset_index_offset().is_some());
5076        assert!(row_group.column(0).column_index_offset().is_some());
5077        // Column "b" should only have offset index
5078        assert!(row_group.column(1).offset_index_offset().is_some());
5079        assert!(row_group.column(1).column_index_offset().is_none());
5080
5081        let options = ReadOptionsBuilder::new().with_page_index().build();
5082        let reader = SerializedFileReader::new_with_options(Bytes::from(buf), options).unwrap();
5083
5084        let row_group = reader.get_row_group(0).unwrap();
5085        let a_col = row_group.metadata().column(0);
5086        let b_col = row_group.metadata().column(1);
5087
5088        // Column chunk of column "a" should have chunk level statistics
5089        if let Statistics::ByteArray(byte_array_stats) = a_col.statistics().unwrap() {
5090            let min = byte_array_stats.min_opt().unwrap();
5091            let max = byte_array_stats.max_opt().unwrap();
5092
5093            assert_eq!(min.as_bytes(), b"a");
5094            assert_eq!(max.as_bytes(), b"d");
5095        } else {
5096            panic!("expecting Statistics::ByteArray");
5097        }
5098
5099        // The column chunk for column "b" shouldn't have statistics
5100        assert!(b_col.statistics().is_none());
5101
5102        let page_index = reader.metadata().page_index().unwrap();
5103
5104        let a_idx = page_index.column_index(0, 0);
5105        assert!(
5106            matches!(a_idx, Some(ColumnIndexMetaData::BYTE_ARRAY(_))),
5107            "{a_idx:?}"
5108        );
5109        let b_idx = page_index.column_index(0, 1);
5110        assert!(b_idx.is_none(), "{b_idx:?}");
5111    }
5112
5113    #[test]
5114    fn test_disabled_statistics_with_chunk() {
5115        let file_schema = Schema::new(vec![
5116            Field::new("a", DataType::Utf8, true),
5117            Field::new("b", DataType::Utf8, true),
5118        ]);
5119        let file_schema = Arc::new(file_schema);
5120
5121        let batch = RecordBatch::try_new(
5122            file_schema.clone(),
5123            vec![
5124                Arc::new(StringArray::from(vec!["a", "b", "c", "d"])) as _,
5125                Arc::new(StringArray::from(vec!["w", "x", "y", "z"])) as _,
5126            ],
5127        )
5128        .unwrap();
5129
5130        let props = WriterProperties::builder()
5131            .set_statistics_enabled(EnabledStatistics::None)
5132            .set_column_statistics_enabled("a".into(), EnabledStatistics::Chunk)
5133            .build();
5134
5135        let mut buf = Vec::with_capacity(1024);
5136        let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), Some(props)).unwrap();
5137        writer.write(&batch).unwrap();
5138
5139        let metadata = writer.close().unwrap();
5140        assert_eq!(metadata.num_row_groups(), 1);
5141        let row_group = metadata.row_group(0);
5142        assert_eq!(row_group.num_columns(), 2);
5143        // Column "a" should only have offset index
5144        assert!(row_group.column(0).offset_index_offset().is_some());
5145        assert!(row_group.column(0).column_index_offset().is_none());
5146        // Column "b" should only have offset index
5147        assert!(row_group.column(1).offset_index_offset().is_some());
5148        assert!(row_group.column(1).column_index_offset().is_none());
5149
5150        let options = ReadOptionsBuilder::new().with_page_index().build();
5151        let reader = SerializedFileReader::new_with_options(Bytes::from(buf), options).unwrap();
5152
5153        let row_group = reader.get_row_group(0).unwrap();
5154        let a_col = row_group.metadata().column(0);
5155        let b_col = row_group.metadata().column(1);
5156
5157        // Column chunk of column "a" should have chunk level statistics
5158        if let Statistics::ByteArray(byte_array_stats) = a_col.statistics().unwrap() {
5159            let min = byte_array_stats.min_opt().unwrap();
5160            let max = byte_array_stats.max_opt().unwrap();
5161
5162            assert_eq!(min.as_bytes(), b"a");
5163            assert_eq!(max.as_bytes(), b"d");
5164        } else {
5165            panic!("expecting Statistics::ByteArray");
5166        }
5167
5168        // The column chunk for column "b"  shouldn't have statistics
5169        assert!(b_col.statistics().is_none());
5170
5171        let page_index = reader.metadata().page_index().unwrap();
5172
5173        let a_idx = page_index.column_index(0, 0);
5174        assert!(a_idx.is_none(), "{a_idx:?}");
5175        let b_idx = page_index.column_index(0, 1);
5176        assert!(b_idx.is_none(), "{b_idx:?}");
5177    }
5178
5179    #[test]
5180    fn test_arrow_writer_skip_metadata() {
5181        let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
5182        let file_schema = Arc::new(batch_schema.clone());
5183
5184        let batch = RecordBatch::try_new(
5185            Arc::new(batch_schema),
5186            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5187        )
5188        .unwrap();
5189        let skip_options = ArrowWriterOptions::new().with_skip_arrow_metadata(true);
5190
5191        let mut buf = Vec::with_capacity(1024);
5192        let mut writer =
5193            ArrowWriter::try_new_with_options(&mut buf, file_schema.clone(), skip_options).unwrap();
5194        writer.write(&batch).unwrap();
5195        writer.close().unwrap();
5196
5197        let bytes = Bytes::from(buf);
5198        let reader_builder = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap();
5199        assert_eq!(file_schema, *reader_builder.schema());
5200        if let Some(key_value_metadata) = reader_builder
5201            .metadata()
5202            .file_metadata()
5203            .key_value_metadata()
5204        {
5205            assert!(
5206                !key_value_metadata
5207                    .iter()
5208                    .any(|kv| kv.key.as_str() == ARROW_SCHEMA_META_KEY)
5209            );
5210        }
5211    }
5212
5213    #[test]
5214    fn test_arrow_writer_skip_path_in_schema() {
5215        let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
5216        let file_schema = Arc::new(batch_schema.clone());
5217
5218        let batch = RecordBatch::try_new(
5219            Arc::new(batch_schema),
5220            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5221        )
5222        .unwrap();
5223
5224        // default options should still write path_in_schema
5225        let skip_options = ArrowWriterOptions::new();
5226
5227        let mut buf = Vec::with_capacity(1024);
5228        let mut writer =
5229            ArrowWriter::try_new_with_options(&mut buf, file_schema.clone(), skip_options).unwrap();
5230        writer.write(&batch).unwrap();
5231        writer.close().unwrap();
5232
5233        // override to not write path_in_schema
5234        let skip_options = ArrowWriterOptions::new().with_properties(
5235            WriterProperties::builder()
5236                .set_write_path_in_schema(false)
5237                .build(),
5238        );
5239
5240        let mut buf2 = Vec::with_capacity(1024);
5241        let mut writer =
5242            ArrowWriter::try_new_with_options(&mut buf2, file_schema.clone(), skip_options)
5243                .unwrap();
5244        writer.write(&batch).unwrap();
5245        writer.close().unwrap();
5246
5247        // buf2 should be a bit smaller due to lack of path_in_schema
5248        assert!(buf.len() > buf2.len());
5249    }
5250
5251    #[test]
5252    fn mismatched_schemas() {
5253        let batch_schema = Schema::new(vec![Field::new("count", DataType::Int32, false)]);
5254        let file_schema = Arc::new(Schema::new(vec![Field::new(
5255            "temperature",
5256            DataType::Float64,
5257            false,
5258        )]));
5259
5260        let batch = RecordBatch::try_new(
5261            Arc::new(batch_schema),
5262            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5263        )
5264        .unwrap();
5265
5266        let mut buf = Vec::with_capacity(1024);
5267        let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), None).unwrap();
5268
5269        let err = writer.write(&batch).unwrap_err().to_string();
5270        assert_eq!(
5271            err,
5272            "Arrow: Incompatible type. Field 'temperature' has type Float64, array has type Int32"
5273        );
5274    }
5275
5276    #[test]
5277    fn test_page_stats_not_written_by_default() {
5278        let string_field = Field::new("a", DataType::Utf8, false);
5279        let schema = Schema::new(vec![string_field]);
5280        let raw_string_values = vec!["Blart Versenwald III"];
5281        let string_values = StringArray::from(raw_string_values.clone());
5282        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(string_values)]).unwrap();
5283
5284        let props = WriterProperties::builder()
5285            .set_statistics_enabled(EnabledStatistics::Page)
5286            .set_dictionary_enabled(false)
5287            .set_encoding(Encoding::PLAIN)
5288            .set_compression(crate::basic::Compression::UNCOMPRESSED)
5289            .build();
5290
5291        let file = roundtrip_opts(&batch, props);
5292
5293        // read file and decode page headers
5294        // Note: use the thrift API as there is no Rust API to access the statistics in the page headers
5295
5296        // decode first page header
5297        let first_page = &file[4..];
5298        let mut prot = ThriftSliceInputProtocol::new(first_page);
5299        let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5300        let stats = hdr.data_page_header.unwrap().statistics;
5301
5302        assert!(stats.is_none());
5303    }
5304
5305    #[test]
5306    fn test_page_stats_when_enabled() {
5307        let string_field = Field::new("a", DataType::Utf8, false);
5308        let schema = Schema::new(vec![string_field]);
5309        let raw_string_values = vec!["Blart Versenwald III", "Andrew Lamb"];
5310        let string_values = StringArray::from(raw_string_values.clone());
5311        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(string_values)]).unwrap();
5312
5313        let props = WriterProperties::builder()
5314            .set_statistics_enabled(EnabledStatistics::Page)
5315            .set_dictionary_enabled(false)
5316            .set_encoding(Encoding::PLAIN)
5317            .set_write_page_header_statistics(true)
5318            .set_compression(crate::basic::Compression::UNCOMPRESSED)
5319            .build();
5320
5321        let file = roundtrip_opts(&batch, props);
5322
5323        // read file and decode page headers
5324        // Note: use the thrift API as there is no Rust API to access the statistics in the page headers
5325
5326        // decode first page header
5327        let first_page = &file[4..];
5328        let mut prot = ThriftSliceInputProtocol::new(first_page);
5329        let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5330        let stats = hdr.data_page_header.unwrap().statistics;
5331
5332        let stats = stats.unwrap();
5333        // check that min/max were actually written to the page
5334        assert!(stats.is_max_value_exact.unwrap());
5335        assert!(stats.is_min_value_exact.unwrap());
5336        assert_eq!(stats.max_value.unwrap(), b"Blart Versenwald III");
5337        assert_eq!(stats.min_value.unwrap(), b"Andrew Lamb");
5338    }
5339
5340    #[test]
5341    fn test_page_stats_truncation() {
5342        let string_field = Field::new("a", DataType::Utf8, false);
5343        let binary_field = Field::new("b", DataType::Binary, false);
5344        let schema = Schema::new(vec![string_field, binary_field]);
5345
5346        let raw_string_values = vec!["Blart Versenwald III"];
5347        let raw_binary_values = [b"Blart Versenwald III".to_vec()];
5348        let raw_binary_value_refs = raw_binary_values
5349            .iter()
5350            .map(|x| x.as_slice())
5351            .collect::<Vec<_>>();
5352
5353        let string_values = StringArray::from(raw_string_values.clone());
5354        let binary_values = BinaryArray::from(raw_binary_value_refs);
5355        let batch = RecordBatch::try_new(
5356            Arc::new(schema),
5357            vec![Arc::new(string_values), Arc::new(binary_values)],
5358        )
5359        .unwrap();
5360
5361        let props = WriterProperties::builder()
5362            .set_statistics_truncate_length(Some(2))
5363            .set_dictionary_enabled(false)
5364            .set_encoding(Encoding::PLAIN)
5365            .set_write_page_header_statistics(true)
5366            .set_compression(crate::basic::Compression::UNCOMPRESSED)
5367            .build();
5368
5369        let file = roundtrip_opts(&batch, props);
5370
5371        // read file and decode page headers
5372        // Note: use the thrift API as there is no Rust API to access the statistics in the page headers
5373
5374        // decode first page header
5375        let first_page = &file[4..];
5376        let mut prot = ThriftSliceInputProtocol::new(first_page);
5377        let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5378        let stats = hdr.data_page_header.unwrap().statistics;
5379        assert!(stats.is_some());
5380        let stats = stats.unwrap();
5381        // check that min/max were properly truncated
5382        assert!(!stats.is_max_value_exact.unwrap());
5383        assert!(!stats.is_min_value_exact.unwrap());
5384        assert_eq!(stats.max_value.unwrap(), b"Bm");
5385        assert_eq!(stats.min_value.unwrap(), b"Bl");
5386
5387        // check second page now
5388        let second_page = &prot.as_slice()[hdr.compressed_page_size as usize..];
5389        let mut prot = ThriftSliceInputProtocol::new(second_page);
5390        let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5391        let stats = hdr.data_page_header.unwrap().statistics;
5392        assert!(stats.is_some());
5393        let stats = stats.unwrap();
5394        // check that min/max were properly truncated
5395        assert!(!stats.is_max_value_exact.unwrap());
5396        assert!(!stats.is_min_value_exact.unwrap());
5397        assert_eq!(stats.max_value.unwrap(), b"Bm");
5398        assert_eq!(stats.min_value.unwrap(), b"Bl");
5399    }
5400
5401    #[test]
5402    fn test_page_encoding_statistics_roundtrip() {
5403        let batch_schema = Schema::new(vec![Field::new(
5404            "int32",
5405            arrow_schema::DataType::Int32,
5406            false,
5407        )]);
5408
5409        let batch = RecordBatch::try_new(
5410            Arc::new(batch_schema.clone()),
5411            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5412        )
5413        .unwrap();
5414
5415        let mut file: File = tempfile::tempfile().unwrap();
5416        let mut writer = ArrowWriter::try_new(&mut file, Arc::new(batch_schema), None).unwrap();
5417        writer.write(&batch).unwrap();
5418        let file_metadata = writer.close().unwrap();
5419
5420        assert_eq!(file_metadata.num_row_groups(), 1);
5421        assert_eq!(file_metadata.row_group(0).num_columns(), 1);
5422        assert!(
5423            file_metadata
5424                .row_group(0)
5425                .column(0)
5426                .page_encoding_stats()
5427                .is_some()
5428        );
5429        let chunk_page_stats = file_metadata
5430            .row_group(0)
5431            .column(0)
5432            .page_encoding_stats()
5433            .unwrap();
5434
5435        // check that the read metadata is also correct
5436        let options = ReadOptionsBuilder::new()
5437            .with_page_index()
5438            .with_encoding_stats_as_mask(false)
5439            .build();
5440        let reader = SerializedFileReader::new_with_options(file, options).unwrap();
5441
5442        let rowgroup = reader.get_row_group(0).expect("row group missing");
5443        assert_eq!(rowgroup.num_columns(), 1);
5444        let column = rowgroup.metadata().column(0);
5445        assert!(column.page_encoding_stats().is_some());
5446        let file_page_stats = column.page_encoding_stats().unwrap();
5447        assert_eq!(chunk_page_stats, file_page_stats);
5448    }
5449
5450    #[test]
5451    #[cfg_attr(miri, ignore)] // Takes too long
5452    fn test_different_dict_page_size_limit() {
5453        let array = Arc::new(Int64Array::from_iter(0..1024 * 1024));
5454        let schema = Arc::new(Schema::new(vec![
5455            Field::new("col0", arrow_schema::DataType::Int64, false),
5456            Field::new("col1", arrow_schema::DataType::Int64, false),
5457        ]));
5458        let batch =
5459            arrow_array::RecordBatch::try_new(schema.clone(), vec![array.clone(), array]).unwrap();
5460
5461        let props = WriterProperties::builder()
5462            .set_dictionary_page_size_limit(1024 * 1024)
5463            .set_column_dictionary_page_size_limit(ColumnPath::from("col1"), 1024 * 1024 * 4)
5464            .build();
5465        let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5466        writer.write(&batch).unwrap();
5467        let data = Bytes::from(writer.into_inner().unwrap());
5468
5469        let mut metadata = ParquetMetaDataReader::new();
5470        metadata.try_parse(&data).unwrap();
5471        let metadata = metadata.finish().unwrap();
5472        let col0_meta = metadata.row_group(0).column(0);
5473        let col1_meta = metadata.row_group(0).column(1);
5474
5475        let get_dict_page_size = move |meta: &ColumnChunkMetaData| {
5476            let mut reader =
5477                SerializedPageReader::new(Arc::new(data.clone()), meta, 0, None).unwrap();
5478            let page = reader.get_next_page().unwrap().unwrap();
5479            match page {
5480                Page::DictionaryPage { buf, .. } => buf.len(),
5481                _ => panic!("expected DictionaryPage"),
5482            }
5483        };
5484
5485        assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
5486        assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
5487    }
5488
5489    #[test]
5490    #[cfg_attr(miri, ignore)] // Takes too long
5491    fn test_arrow_writer_granular_mode_roundtrip() {
5492        // Granular mode subdivides chunks and writes more pages than the
5493        // default batched path. Make sure the data we write back is
5494        // bit-identical to what went in — page-count assertions elsewhere
5495        // only prove pages were cut, not that the encoded data is correct.
5496        //
5497        // Mix value sizes so that the cumulative-byte-budget cutoff
5498        // lands mid-chunk, exercising both batched and granular paths
5499        // within the same `write_batch_internal` call.
5500        let small = "tiny".to_string();
5501        let big = "x".repeat(64 * 1024);
5502        let strings: Vec<String> = (0..256)
5503            .map(|i| {
5504                if i % 16 == 0 {
5505                    big.clone()
5506                } else {
5507                    small.clone()
5508                }
5509            })
5510            .collect();
5511
5512        let schema = Arc::new(Schema::new(vec![Field::new(
5513            "col",
5514            ArrowDataType::Utf8,
5515            false,
5516        )]));
5517        let batch = RecordBatch::try_new(
5518            schema.clone(),
5519            vec![Arc::new(StringArray::from(strings.clone())) as _],
5520        )
5521        .unwrap();
5522
5523        let props = WriterProperties::builder()
5524            .set_dictionary_enabled(false)
5525            .set_data_page_size_limit(16 * 1024)
5526            .build();
5527        let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5528        writer.write(&batch).unwrap();
5529        let data = Bytes::from(writer.into_inner().unwrap());
5530
5531        let mut reader = ParquetRecordBatchReader::try_new(data, 1024).unwrap();
5532        let read = reader.next().unwrap().unwrap();
5533        assert!(reader.next().is_none(), "expected one batch");
5534        let col = read
5535            .column(0)
5536            .as_any()
5537            .downcast_ref::<StringArray>()
5538            .unwrap();
5539        assert_eq!(col.len(), strings.len());
5540        for (i, expected) in strings.iter().enumerate() {
5541            assert_eq!(
5542                col.value(i),
5543                expected.as_str(),
5544                "value mismatch at index {i}"
5545            );
5546        }
5547    }
5548
5549    #[test]
5550    fn test_arrow_writer_all_null_string_column() {
5551        // The `LevelDataRef::value_count` Uniform branch with
5552        // `value != max_def` (entirely-null chunk) must return 0 so the
5553        // sub-batch sizer short-circuits to batch mode without trying
5554        // to estimate byte budgets for non-existent values.
5555        let num_rows = 1024;
5556        let schema = Arc::new(Schema::new(vec![Field::new(
5557            "col",
5558            ArrowDataType::Utf8,
5559            true,
5560        )]));
5561        let nulls: Vec<Option<&str>> = vec![None; num_rows];
5562        let batch = RecordBatch::try_new(
5563            schema.clone(),
5564            vec![Arc::new(StringArray::from(nulls)) as _],
5565        )
5566        .unwrap();
5567
5568        let props = WriterProperties::builder()
5569            .set_dictionary_enabled(false)
5570            .set_data_page_size_limit(16 * 1024)
5571            .build();
5572        let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5573        writer.write(&batch).unwrap();
5574        let data = Bytes::from(writer.into_inner().unwrap());
5575
5576        // Re-parse the file: row group has one column, every row is
5577        // null, all data pages report `num_rows / page_count` rows.
5578        let mut metadata = ParquetMetaDataReader::new();
5579        metadata.try_parse(&data).unwrap();
5580        let metadata = metadata.finish().unwrap();
5581        let row_group = metadata.row_group(0);
5582        let col_meta = row_group.column(0);
5583        assert_eq!(row_group.num_rows() as usize, num_rows);
5584        // Statistics record `null_count = num_rows` — proves every value
5585        // was written as null.
5586        if let Some(stats) = col_meta.statistics() {
5587            assert_eq!(
5588                stats.null_count_opt().unwrap_or(0) as usize,
5589                num_rows,
5590                "expected all-null column to report null_count = num_rows"
5591            );
5592        }
5593
5594        let mut reader =
5595            SerializedPageReader::new(Arc::new(data.clone()), col_meta, num_rows, None).unwrap();
5596        let mut total_values = 0u32;
5597        while let Some(page) = reader.get_next_page().unwrap() {
5598            if matches!(page, Page::DataPage { .. } | Page::DataPageV2 { .. }) {
5599                total_values += page.num_values();
5600            }
5601        }
5602        assert_eq!(
5603            total_values as usize, num_rows,
5604            "expected every level position to be represented in some page"
5605        );
5606    }
5607
5608    struct WriteBatchesShape {
5609        num_batches: usize,
5610        rows_per_batch: usize,
5611        row_size: usize,
5612    }
5613
5614    /// Helper function to write batches with the provided `WriteBatchesShape` into an `ArrowWriter`
5615    fn write_batches(
5616        WriteBatchesShape {
5617            num_batches,
5618            rows_per_batch,
5619            row_size,
5620        }: WriteBatchesShape,
5621        props: WriterProperties,
5622    ) -> ParquetRecordBatchReaderBuilder<File> {
5623        let schema = Arc::new(Schema::new(vec![Field::new(
5624            "str",
5625            ArrowDataType::Utf8,
5626            false,
5627        )]));
5628        let file = tempfile::tempfile().unwrap();
5629        let mut writer =
5630            ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
5631
5632        for batch_idx in 0..num_batches {
5633            let strings: Vec<String> = (0..rows_per_batch)
5634                .map(|i| format!("{:0>width$}", batch_idx * 10 + i, width = row_size))
5635                .collect();
5636            let array = StringArray::from(strings);
5637            let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
5638            writer.write(&batch).unwrap();
5639        }
5640        writer.close().unwrap();
5641        ParquetRecordBatchReaderBuilder::try_new(file).unwrap()
5642    }
5643
5644    #[test]
5645    // When both limits are None, all data should go into a single row group
5646    fn test_row_group_limit_none_writes_single_row_group() {
5647        let props = WriterProperties::builder()
5648            .set_max_row_group_row_count(None)
5649            .set_max_row_group_bytes(None)
5650            .build();
5651
5652        let builder = write_batches(
5653            WriteBatchesShape {
5654                num_batches: 1,
5655                rows_per_batch: 1000,
5656                row_size: 4,
5657            },
5658            props,
5659        );
5660
5661        assert_eq!(
5662            &row_group_sizes(builder.metadata()),
5663            &[1000],
5664            "With no limits, all rows should be in a single row group"
5665        );
5666    }
5667
5668    #[test]
5669    // When only max_row_group_size is set, respect the row limit
5670    fn test_row_group_limit_rows_only() {
5671        let props = WriterProperties::builder()
5672            .set_max_row_group_row_count(Some(300))
5673            .set_max_row_group_bytes(None)
5674            .build();
5675
5676        let builder = write_batches(
5677            WriteBatchesShape {
5678                num_batches: 1,
5679                rows_per_batch: 1000,
5680                row_size: 4,
5681            },
5682            props,
5683        );
5684
5685        assert_eq!(
5686            &row_group_sizes(builder.metadata()),
5687            &[300, 300, 300, 100],
5688            "Row groups should be split by row count"
5689        );
5690    }
5691
5692    #[test]
5693    #[cfg_attr(miri, ignore)] // Takes too long
5694    // A row limit far smaller than the batch splits it many times over; the split must not
5695    // consume stack proportional to the number of row groups.
5696    fn test_row_group_limit_rows_only_many_splits() {
5697        let props = WriterProperties::builder()
5698            .set_max_row_group_row_count(Some(1))
5699            .set_max_row_group_bytes(None)
5700            .build();
5701
5702        let rows = 50_000;
5703        let builder = write_batches(
5704            WriteBatchesShape {
5705                num_batches: 1,
5706                rows_per_batch: rows,
5707                row_size: 4,
5708            },
5709            props,
5710        );
5711
5712        let sizes = row_group_sizes(builder.metadata());
5713        assert_eq!(sizes.len(), rows, "Every row should get its own row group");
5714        assert_eq!(
5715            sizes.iter().sum::<i64>(),
5716            rows as i64,
5717            "Total rows should be preserved"
5718        );
5719    }
5720
5721    #[test]
5722    // When only max_row_group_bytes is set, respect the byte limit
5723    fn test_row_group_limit_bytes_only() {
5724        let props = WriterProperties::builder()
5725            .set_max_row_group_row_count(None)
5726            // Set byte limit to approximately fit ~30 rows worth of data (~100 bytes each)
5727            .set_max_row_group_bytes(Some(3500))
5728            .build();
5729
5730        let builder = write_batches(
5731            WriteBatchesShape {
5732                num_batches: 10,
5733                rows_per_batch: 10,
5734                row_size: 100,
5735            },
5736            props,
5737        );
5738
5739        let sizes = row_group_sizes(builder.metadata());
5740
5741        assert!(
5742            sizes.len() > 1,
5743            "Should have multiple row groups due to byte limit, got {sizes:?}",
5744        );
5745
5746        let total_rows: i64 = sizes.iter().sum();
5747        assert_eq!(total_rows, 100, "Total rows should be preserved");
5748    }
5749
5750    #[test]
5751    // If an in-progress row group is already oversized, it should be flushed before writing more.
5752    fn test_row_group_limit_bytes_flushes_when_current_group_already_too_large() {
5753        let schema = Arc::new(Schema::new(vec![Field::new(
5754            "str",
5755            ArrowDataType::Utf8,
5756            false,
5757        )]));
5758        let file = tempfile::tempfile().unwrap();
5759
5760        // Start with no byte limit so we can intentionally build an oversized in-progress row group.
5761        let props = WriterProperties::builder()
5762            .set_max_row_group_row_count(None)
5763            .set_max_row_group_bytes(None)
5764            .build();
5765        let mut writer =
5766            ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
5767
5768        let first_array = StringArray::from(
5769            (0..10)
5770                .map(|i| format!("{i:0>100}"))
5771                .collect::<Vec<String>>(),
5772        );
5773        let first_batch =
5774            RecordBatch::try_new(schema.clone(), vec![Arc::new(first_array)]).unwrap();
5775        writer.write(&first_batch).unwrap();
5776        assert_eq!(writer.in_progress_rows(), 10);
5777
5778        // Tighten the limit below the current in-progress bytes to exercise:
5779        // `if current_bytes >= max_bytes { self.flush()?; ... }`
5780        writer.max_row_group_bytes = Some(1);
5781
5782        let second_array = StringArray::from(vec!["x".to_string()]);
5783        let second_batch =
5784            RecordBatch::try_new(schema.clone(), vec![Arc::new(second_array)]).unwrap();
5785        writer.write(&second_batch).unwrap();
5786        writer.close().unwrap();
5787        let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
5788
5789        assert_eq!(
5790            &row_group_sizes(builder.metadata()),
5791            &[10, 1],
5792            "The second write should flush an oversized in-progress row group first",
5793        );
5794    }
5795
5796    #[test]
5797    // When both limits are set, the row limit triggers first
5798    fn test_row_group_limit_both_row_wins_single_batch() {
5799        let props = WriterProperties::builder()
5800            .set_max_row_group_row_count(Some(200)) // Will trigger at 200 rows
5801            .set_max_row_group_bytes(Some(1024 * 1024)) // 1MB - won't trigger for small int data
5802            .build();
5803
5804        let builder = write_batches(
5805            WriteBatchesShape {
5806                num_batches: 1,
5807                row_size: 4,
5808                rows_per_batch: 1000,
5809            },
5810            props,
5811        );
5812
5813        assert_eq!(
5814            &row_group_sizes(builder.metadata()),
5815            &[200, 200, 200, 200, 200],
5816            "Row limit should trigger before byte limit"
5817        );
5818    }
5819
5820    #[test]
5821    // When both limits are set, the row limit triggers first
5822    fn test_row_group_limit_both_row_wins_multiple_batches() {
5823        let props = WriterProperties::builder()
5824            .set_max_row_group_row_count(Some(5)) // Will trigger every 5 rows
5825            .set_max_row_group_bytes(Some(9999)) // Won't trigger
5826            .build();
5827
5828        let builder = write_batches(
5829            WriteBatchesShape {
5830                num_batches: 10,
5831                rows_per_batch: 10,
5832                row_size: 100,
5833            },
5834            props,
5835        );
5836
5837        assert_eq!(
5838            &row_group_sizes(builder.metadata()),
5839            &[5; 20],
5840            "Row limit should trigger before byte limit"
5841        );
5842    }
5843
5844    #[test]
5845    // When both limits are set, the byte limit triggers first
5846    fn test_row_group_limit_both_bytes_wins() {
5847        let props = WriterProperties::builder()
5848            .set_max_row_group_row_count(Some(1000)) // Won't trigger for 100 rows
5849            .set_max_row_group_bytes(Some(3500)) // Will trigger at ~30-35 rows
5850            .build();
5851
5852        let builder = write_batches(
5853            WriteBatchesShape {
5854                num_batches: 10,
5855                rows_per_batch: 10,
5856                row_size: 100,
5857            },
5858            props,
5859        );
5860
5861        let sizes = row_group_sizes(builder.metadata());
5862
5863        assert!(
5864            sizes.len() > 1,
5865            "Byte limit should trigger before row limit, got {sizes:?}",
5866        );
5867
5868        assert!(
5869            sizes.iter().all(|&s| s < 1000),
5870            "No row group should hit the row limit"
5871        );
5872
5873        let total_rows: i64 = sizes.iter().sum();
5874        assert_eq!(total_rows, 100, "Total rows should be preserved");
5875    }
5876
5877    #[test]
5878    // Both limits can apply to the same batch: the row limit trims it to 5 rows, and the
5879    // byte limit then trims those 5 down to 4.
5880    fn test_row_group_limit_both_apply_to_same_batch() {
5881        let props = WriterProperties::builder()
5882            .set_max_row_group_row_count(Some(15))
5883            .set_max_row_group_bytes(Some(1500))
5884            .build();
5885
5886        let builder = write_batches(
5887            WriteBatchesShape {
5888                num_batches: 2,
5889                rows_per_batch: 10,
5890                row_size: 100,
5891            },
5892            props,
5893        );
5894
5895        assert_eq!(
5896            &row_group_sizes(builder.metadata()),
5897            &[14, 6],
5898            "Byte limit should still apply to a batch the row limit already split"
5899        );
5900    }
5901
5902    #[test]
5903    // A batch far larger than the byte limit must be split even when it is the first one
5904    // written to a row group, where there are no buffered rows to estimate the row size from.
5905    fn test_row_group_limit_bytes_splits_first_batch() {
5906        let props = WriterProperties::builder()
5907            .set_max_row_group_row_count(None)
5908            .set_max_row_group_bytes(Some(3500))
5909            .build();
5910
5911        let builder = write_batches(
5912            WriteBatchesShape {
5913                num_batches: 3,
5914                rows_per_batch: 100,
5915                row_size: 100,
5916            },
5917            props,
5918        );
5919
5920        let sizes = row_group_sizes(builder.metadata());
5921        assert!(
5922            sizes.iter().all(|&s| s < 100),
5923            "No row group should hold a whole batch, got {sizes:?}",
5924        );
5925        assert_eq!(
5926            sizes.iter().sum::<i64>(),
5927            300,
5928            "Total rows should be preserved"
5929        );
5930        for row_group in builder.metadata().row_groups() {
5931            assert!(
5932                row_group.total_byte_size() <= 3500 + 100,
5933                "Row groups should stay within one row of the byte limit, got {}",
5934                row_group.total_byte_size(),
5935            );
5936        }
5937    }
5938
5939    #[test]
5940    fn arrow_column_chunk_close_mut_drops_column_index() {
5941        use crate::arrow::ArrowSchemaConverter;
5942        use crate::file::writer::SerializedFileWriter;
5943
5944        let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)]));
5945        let props = Arc::new(
5946            WriterProperties::builder()
5947                .set_statistics_enabled(EnabledStatistics::Page)
5948                .build(),
5949        );
5950        let parquet_schema = ArrowSchemaConverter::new()
5951            .with_coerce_types(props.coerce_types())
5952            .convert(&schema)
5953            .unwrap();
5954
5955        let mut buf = Vec::with_capacity(1024);
5956        let mut writer =
5957            SerializedFileWriter::new(&mut buf, parquet_schema.root_schema_ptr(), props.clone())
5958                .unwrap();
5959
5960        let factory = ArrowRowGroupWriterFactory::try_new(&writer, Arc::clone(&schema)).unwrap();
5961        let mut col_writers = factory.create_column_writers(0).unwrap();
5962        let arr: ArrayRef = Arc::new(Int32Array::from_iter_values(0..64));
5963        for leaves in compute_leaves(schema.field(0), &arr).unwrap() {
5964            col_writers[0].write(&leaves).unwrap();
5965        }
5966        let mut chunk = col_writers.pop().unwrap().close().unwrap();
5967
5968        // Immutable accessor exposes the close result produced at close time.
5969        assert!(
5970            chunk.close().column_index.is_some(),
5971            "EnabledStatistics::Page should produce a column_index"
5972        );
5973
5974        // Mutable accessor lets callers drop the page-level index before append.
5975        chunk.close_mut().column_index = None;
5976        assert!(chunk.close().column_index.is_none());
5977
5978        let mut rg = writer.next_row_group().unwrap();
5979        chunk.append_to_row_group(&mut rg).unwrap();
5980        rg.close().unwrap();
5981        let file_meta = writer.close().unwrap();
5982
5983        // After dropping column_index, the resulting file records no column
5984        // index offset/length for this chunk.
5985        let cc = file_meta.row_group(0).column(0);
5986        assert!(cc.column_index_range().is_none());
5987    }
5988
5989    #[test]
5990    #[cfg_attr(miri, ignore)] // Takes too long
5991    fn test_number_distinct_values_exact_count() {
5992        // 50 distinct Int32 values repeated across 100k rows, with every 7th row null.
5993        // Nulls must not be counted as a distinct value.
5994        let cardinality = 50u32;
5995        let array: ArrayRef = Arc::new(Int32Array::from_iter((0..100_000u32).map(|i| {
5996            if i % 7 == 0 {
5997                None
5998            } else {
5999                Some((i % cardinality) as i32)
6000            }
6001        })));
6002        let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, true)]));
6003        let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
6004
6005        let props = WriterProperties::builder()
6006            .set_write_row_group_number_distinct_values(true)
6007            .build();
6008        let mut buf = Vec::new();
6009        let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), Some(props)).unwrap();
6010        writer.write(&batch).unwrap();
6011        let metadata = writer.close().unwrap();
6012
6013        let count = metadata
6014            .row_group(0)
6015            .column(0)
6016            .statistics()
6017            .and_then(|s| s.distinct_count_opt())
6018            .expect("distinct_count should be set");
6019        // Must equal cardinality exactly; nulls must not inflate the count.
6020        assert_eq!(count, cardinality as u64);
6021    }
6022
6023    #[test]
6024    fn test_number_distinct_values_view_types() {
6025        // 5 distinct values repeated across 30 rows, with every 4th row null.
6026        // Verifies Utf8View is counted correctly (BinaryView shares the same code path).
6027        let cardinality = 5u32;
6028        let distinct_strings = ["alpha", "beta", "gamma", "delta", "epsilon"];
6029
6030        let string_view_col: ArrayRef = Arc::new(StringViewArray::from_iter((0..30u32).map(|i| {
6031            if i % 4 == 0 {
6032                None
6033            } else {
6034                Some(distinct_strings[(i % cardinality) as usize])
6035            }
6036        })));
6037
6038        let schema = Arc::new(Schema::new(vec![Field::new(
6039            "string_view_col",
6040            DataType::Utf8View,
6041            true,
6042        )]));
6043        let batch = RecordBatch::try_new(schema, vec![string_view_col]).unwrap();
6044
6045        let props = WriterProperties::builder()
6046            .set_write_row_group_number_distinct_values(true)
6047            .build();
6048        let mut parquet_bytes = Vec::new();
6049        let mut writer =
6050            ArrowWriter::try_new(&mut parquet_bytes, batch.schema(), Some(props)).unwrap();
6051        writer.write(&batch).unwrap();
6052        let metadata = writer.close().unwrap();
6053
6054        let distinct_count = metadata
6055            .row_group(0)
6056            .column(0)
6057            .statistics()
6058            .and_then(|s| s.distinct_count_opt())
6059            .expect("distinct_count should be set for Utf8View column");
6060        assert_eq!(distinct_count, cardinality as u64);
6061    }
6062
6063    #[test]
6064    fn test_number_distinct_values_not_written_by_default() {
6065        let array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..100));
6066        let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, false)]));
6067        let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
6068
6069        let mut buf = Vec::new();
6070        let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), None).unwrap();
6071        writer.write(&batch).unwrap();
6072        let metadata = writer.close().unwrap();
6073
6074        let count = metadata
6075            .row_group(0)
6076            .column(0)
6077            .statistics()
6078            .and_then(|s| s.distinct_count_opt());
6079        assert!(count.is_none());
6080    }
6081
6082    #[test]
6083    fn test_dictionary_ndv_single_batch() {
6084        // Dictionary array with 3 distinct string values repeated many times.
6085        // NDV must equal the number of distinct values in the dictionary (3),
6086        // not the number of rows.
6087        let keys = Int32Array::from(vec![0, 1, 2, 0, 1, 2, 0, 1, 2]);
6088        let values: ArrayRef = Arc::new(StringArray::from(vec!["cat", "dog", "bird"]));
6089        let dict: ArrayRef = Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap());
6090
6091        let schema = Arc::new(Schema::new(vec![Field::new(
6092            "x",
6093            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
6094            false,
6095        )]));
6096        let batch = RecordBatch::try_new(schema, vec![dict]).unwrap();
6097
6098        let props = WriterProperties::builder()
6099            .set_write_row_group_number_distinct_values(true)
6100            .build();
6101        let mut buf = Vec::new();
6102        let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), Some(props)).unwrap();
6103        writer.write(&batch).unwrap();
6104        let metadata = writer.close().unwrap();
6105
6106        let count = metadata
6107            .row_group(0)
6108            .column(0)
6109            .statistics()
6110            .and_then(|s| s.distinct_count_opt())
6111            .expect("distinct_count should be set");
6112        assert_eq!(count, 3);
6113    }
6114
6115    #[test]
6116    fn test_dictionary_ndv_excludes_unreferenced_values() {
6117        // Keys only reference indices 0 and 1; value at index 2 ("unreferenced") should not
6118        // count toward NDV even though it appears in the dictionary's values array.
6119        let keys = Int32Array::from(vec![0, 1, 0, 1]);
6120        let values: ArrayRef = Arc::new(StringArray::from(vec!["cat", "dog", "unreferenced"]));
6121        let dict: ArrayRef = Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap());
6122
6123        let schema = Arc::new(Schema::new(vec![Field::new(
6124            "x",
6125            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
6126            false,
6127        )]));
6128        let batch = RecordBatch::try_new(schema, vec![dict]).unwrap();
6129
6130        let props = WriterProperties::builder()
6131            .set_write_row_group_number_distinct_values(true)
6132            .build();
6133        let mut buf = Vec::new();
6134        let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), Some(props)).unwrap();
6135        writer.write(&batch).unwrap();
6136        let metadata = writer.close().unwrap();
6137
6138        let count = metadata
6139            .row_group(0)
6140            .column(0)
6141            .statistics()
6142            .and_then(|s| s.distinct_count_opt())
6143            .expect("distinct_count should be set");
6144        assert_eq!(
6145            count, 2,
6146            "unreferenced dictionary values must not count toward NDV"
6147        );
6148    }
6149
6150    #[test]
6151    fn test_dictionary_ndv_across_batches_regression() {
6152        // Regression test for https://github.com/apache/arrow-rs/issues/11172.
6153        let make_dict_batch = |a: &str, b: &str| -> RecordBatch {
6154            let keys = Int32Array::from(vec![0, 1, 0, 1]);
6155            let values: ArrayRef = Arc::new(StringArray::from(vec![a, b]));
6156            let dict: ArrayRef =
6157                Arc::new(DictionaryArray::<Int32Type>::try_new(keys, values).unwrap());
6158            let schema = Arc::new(Schema::new(vec![Field::new(
6159                "x",
6160                DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
6161                false,
6162            )]));
6163            RecordBatch::try_new(schema, vec![dict]).unwrap()
6164        };
6165
6166        // batch1: dict = ["cat", "dog"], batch2: dict = ["fish", "cat"]
6167        // Distinct values across both batches: "cat", "dog", "fish" NDV = 3
6168        let batch1 = make_dict_batch("cat", "dog");
6169        let batch2 = make_dict_batch("fish", "cat");
6170
6171        let props = WriterProperties::builder()
6172            .set_write_row_group_number_distinct_values(true)
6173            .build();
6174        let mut buf = Vec::new();
6175        let mut writer = ArrowWriter::try_new(&mut buf, batch1.schema(), Some(props)).unwrap();
6176        writer.write(&batch1).unwrap();
6177        writer.write(&batch2).unwrap();
6178        let metadata = writer.close().unwrap();
6179
6180        let count = metadata
6181            .row_group(0)
6182            .column(0)
6183            .statistics()
6184            .and_then(|s| s.distinct_count_opt())
6185            .expect("distinct_count should be set");
6186        assert_eq!(
6187            count, 3,
6188            "NDV should count distinct values, not distinct key indices"
6189        );
6190    }
6191
6192    #[test]
6193    fn ree_struct_with_ree_child() {
6194        // Struct with a REE string field and a REE int field — confirms
6195        // recursion visits every child and each collapses to the right leaf type.
6196        let run_ends = Int32Array::from(vec![2i32, 3, 5]);
6197
6198        let col_a: ArrayRef = Arc::new(
6199            RunArray::try_new(
6200                &run_ends,
6201                &StringArray::from(vec![Some("foo"), None, Some("bar")]),
6202            )
6203            .unwrap(),
6204        );
6205        let col_b: ArrayRef = Arc::new(
6206            RunArray::try_new(&run_ends, &Int32Array::from(vec![Some(1), None, Some(2)])).unwrap(),
6207        );
6208
6209        let struct_array: ArrayRef = Arc::new(StructArray::new(
6210            Fields::from(vec![
6211                Field::new("a", col_a.data_type().clone(), true),
6212                Field::new("b", col_b.data_type().clone(), true),
6213            ]),
6214            vec![col_a, col_b],
6215            None,
6216        ));
6217
6218        let schema = Arc::new(Schema::new(vec![Field::new(
6219            "row",
6220            struct_array.data_type().clone(),
6221            true,
6222        )]));
6223        let batch = RecordBatch::try_new(schema.clone(), vec![struct_array]).unwrap();
6224
6225        let mut buf = Vec::new();
6226        let mut writer = ArrowWriter::try_new(&mut buf, schema, None).unwrap();
6227        writer.write(&batch).unwrap();
6228        let metadata = writer.close().unwrap();
6229
6230        let parquet_schema = metadata.file_metadata().schema_descr();
6231        assert_eq!(parquet_schema.num_columns(), 2);
6232        assert_eq!(
6233            parquet_schema.column(0).physical_type(),
6234            crate::basic::Type::BYTE_ARRAY
6235        );
6236        assert_eq!(parquet_schema.column(0).path().string(), "row.a");
6237        assert_eq!(
6238            parquet_schema.column(1).physical_type(),
6239            crate::basic::Type::INT32
6240        );
6241        assert_eq!(parquet_schema.column(1).path().string(), "row.b");
6242    }
6243
6244    /// Regression test for https://github.com/apache/arrow-rs/issues/11073
6245    ///
6246    /// This checks that the `ArrowColumnWriter` path produces the same min/max
6247    /// statistics as writing the column directly with the low-level
6248    /// `SerializedFileWriter` API, for both `BinaryArray` (the buggy case)
6249    /// and `FixedSizeBinaryArray` (which already worked correctly).
6250    #[test]
6251    fn test_decimal_byte_array_statistics_signed_comparison() {
6252        use crate::basic::{LogicalType, Repetition, Type as PhysicalType};
6253        use crate::data_type::{ByteArray, ByteArrayType};
6254        use crate::schema::types::Type as SchemaType;
6255
6256        // 1-byte two's-complement decimal values: -1, 0, 1
6257        const VALUES: [&[u8]; 3] = [&[0xffu8], &[0x00], &[0x01]];
6258
6259        fn new_file(physical: PhysicalType) -> SerializedFileWriter<Vec<u8>> {
6260            let mut field = SchemaType::primitive_type_builder("value", physical)
6261                .with_repetition(Repetition::REQUIRED)
6262                .with_logical_type(Some(LogicalType::decimal(0, 2)))
6263                .with_precision(2)
6264                .with_scale(0);
6265            if physical == PhysicalType::FIXED_LEN_BYTE_ARRAY {
6266                field = field.with_length(1);
6267            }
6268            let root = SchemaType::group_type_builder("root")
6269                .with_fields(vec![Arc::new(field.build().unwrap())])
6270                .build()
6271                .unwrap();
6272            SerializedFileWriter::new(Vec::new(), Arc::new(root), Default::default()).unwrap()
6273        }
6274
6275        fn read_minmax(file: Vec<u8>) -> (i8, i8) {
6276            let reader = SerializedFileReader::new(Bytes::from(file)).unwrap();
6277            let metadata = reader.metadata();
6278            assert_eq!(metadata.file_metadata().num_rows(), 3);
6279            let stats = metadata
6280                .row_group(0)
6281                .column(0)
6282                .statistics()
6283                .expect("missing statistics");
6284            (
6285                stats.min_bytes_opt().unwrap()[0] as i8,
6286                stats.max_bytes_opt().unwrap()[0] as i8,
6287            )
6288        }
6289
6290        fn direct_parquet() -> Vec<u8> {
6291            let mut file = new_file(PhysicalType::BYTE_ARRAY);
6292            let mut group = file.next_row_group().unwrap();
6293            let mut column = group.next_column().unwrap().unwrap();
6294            let values = VALUES
6295                .iter()
6296                .map(|v| ByteArray::from(v.to_vec()))
6297                .collect::<Vec<_>>();
6298            column
6299                .typed::<ByteArrayType>()
6300                .write_batch(&values, None, None)
6301                .unwrap();
6302            column.close().unwrap();
6303            group.close().unwrap();
6304            file.into_inner().unwrap()
6305        }
6306
6307        fn arrow_bounds(mut file: SerializedFileWriter<Vec<u8>>, array: ArrayRef) -> Vec<u8> {
6308            let field = Field::new("value", array.data_type().clone(), false);
6309            let schema = Arc::new(Schema::new(vec![field.clone()]));
6310            let factory = ArrowRowGroupWriterFactory::try_new(&file, schema).unwrap();
6311            let mut writer: ArrowColumnWriter = factory.create_column_writers(0).unwrap().remove(0);
6312            let leaf = compute_leaves(&field, &array).unwrap().remove(0);
6313            writer.write(&leaf).unwrap();
6314            let mut group = file.next_row_group().unwrap();
6315            writer
6316                .close()
6317                .unwrap()
6318                .append_to_row_group(&mut group)
6319                .unwrap();
6320            group.close().unwrap();
6321            file.into_inner().unwrap()
6322        }
6323
6324        fn arrow_binary() -> Vec<u8> {
6325            let file = new_file(PhysicalType::BYTE_ARRAY);
6326            let array: ArrayRef = Arc::new(BinaryArray::from_vec(VALUES.to_vec()));
6327            arrow_bounds(file, array)
6328        }
6329
6330        fn arrow_fixed_size_binary() -> Vec<u8> {
6331            let file = new_file(PhysicalType::FIXED_LEN_BYTE_ARRAY);
6332            let array: ArrayRef =
6333                Arc::new(FixedSizeBinaryArray::try_from_iter(VALUES.into_iter()).unwrap());
6334            arrow_bounds(file, array)
6335        }
6336
6337        let direct = read_minmax(direct_parquet());
6338        let fixed = read_minmax(arrow_fixed_size_binary());
6339        let binary = read_minmax(arrow_binary());
6340
6341        assert_eq!(direct, (-1, 1), "incorrect direct Parquet stats");
6342        assert_eq!(fixed, (-1, 1), "incorrect Arrow FixedSizeBinaryArray stats");
6343        assert_eq!(binary, (-1, 1), "incorrect Arrow BinaryArray stats");
6344    }
6345}