Skip to main content

parquet/arrow/arrow_writer/
mod.rs

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