Skip to main content

parquet/arrow/arrow_writer/
mod.rs

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