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    ///
1190    /// # Errors
1191    ///
1192    /// Returns an error if the column could not be finalised, or if another thread
1193    /// panicked while holding the column chunk. The caller cannot cause either.
1194    pub fn close(self) -> Result<ArrowColumnChunk> {
1195        let distinct_count = self
1196            .distinct_values_seen
1197            .as_ref()
1198            .filter(|s| !s.is_empty())
1199            .map(|s| s.len() as u64);
1200        let close = match self.writer {
1201            ArrowColumnWriterImpl::ByteArray(mut c) => {
1202                if let Some(count) = distinct_count {
1203                    c.set_distinct_count_override(count);
1204                }
1205                c.close()?
1206            }
1207            ArrowColumnWriterImpl::Column(mut c) => {
1208                if let Some(count) = distinct_count {
1209                    c.set_distinct_count_override(count);
1210                }
1211                c.close()?
1212            }
1213        };
1214        // Closing the writer above dropped the only other handle on the chunk.
1215        let chunk = Arc::try_unwrap(self.chunk)
1216            .map_err(|_| general_err!("Internal Error: the column chunk is still shared"))?;
1217        let data = chunk
1218            .into_inner()
1219            .map_err(|_| general_err!("The column chunk lock is poisoned"))?;
1220        Ok(ArrowColumnChunk { data, close })
1221    }
1222
1223    /// Returns the estimated total memory usage by the writer.
1224    ///
1225    /// This  [`Self::get_estimated_total_bytes`] this is an estimate
1226    /// of the current memory usage and not it's anticipated encoded size.
1227    ///
1228    /// This includes:
1229    /// 1. Data buffered in encoded form
1230    /// 2. Data buffered in un-encoded form (e.g. `usize` dictionary keys)
1231    ///
1232    /// This value should be greater than or equal to [`Self::get_estimated_total_bytes`]
1233    pub fn memory_size(&self) -> usize {
1234        match &self.writer {
1235            ArrowColumnWriterImpl::ByteArray(c) => c.memory_size(),
1236            ArrowColumnWriterImpl::Column(c) => c.memory_size(),
1237        }
1238    }
1239
1240    /// Returns the estimated total encoded bytes for this column writer.
1241    ///
1242    /// This includes:
1243    /// 1. Data buffered in encoded form
1244    /// 2. An estimate of how large the data buffered in un-encoded form would be once encoded
1245    ///
1246    /// This value should be less than or equal to [`Self::memory_size`]
1247    pub fn get_estimated_total_bytes(&self) -> usize {
1248        match &self.writer {
1249            ArrowColumnWriterImpl::ByteArray(c) => c.get_estimated_total_bytes() as _,
1250            ArrowColumnWriterImpl::Column(c) => c.get_estimated_total_bytes() as _,
1251        }
1252    }
1253}
1254
1255/// Encodes [`RecordBatch`] to a parquet row group
1256///
1257/// Note: this structure is created by [`ArrowRowGroupWriterFactory`] internally used to
1258/// create [`ArrowRowGroupWriter`]s, but it is not exposed publicly.
1259///
1260/// See the example on [`ArrowColumnWriter`] for how to encode columns in parallel
1261#[derive(Debug)]
1262struct ArrowRowGroupWriter {
1263    writers: Vec<ArrowColumnWriter>,
1264    schema: SchemaRef,
1265    buffered_rows: usize,
1266}
1267
1268impl ArrowRowGroupWriter {
1269    fn new(writers: Vec<ArrowColumnWriter>, arrow: &SchemaRef) -> Self {
1270        Self {
1271            writers,
1272            schema: arrow.clone(),
1273            buffered_rows: 0,
1274        }
1275    }
1276
1277    fn write(&mut self, batch: &RecordBatch) -> Result<()> {
1278        self.buffered_rows += batch.num_rows();
1279        let mut writers = self.writers.iter_mut();
1280        for (field, column) in self.schema.fields().iter().zip(batch.columns()) {
1281            for leaf in compute_leaves(field.as_ref(), column)? {
1282                writers.next().unwrap().write(&leaf)?;
1283            }
1284        }
1285        Ok(())
1286    }
1287
1288    fn write_with_chunkers(
1289        &mut self,
1290        batch: &RecordBatch,
1291        chunkers: &mut [ContentDefinedChunker],
1292    ) -> Result<()> {
1293        self.buffered_rows += batch.num_rows();
1294        let mut writers = self.writers.iter_mut();
1295        let mut chunkers = chunkers.iter_mut();
1296        for (field, column) in self.schema.fields().iter().zip(batch.columns()) {
1297            for leaf in compute_leaves(field.as_ref(), column)? {
1298                writers
1299                    .next()
1300                    .unwrap()
1301                    .write_with_chunker(&leaf, chunkers.next().unwrap())?;
1302            }
1303        }
1304        Ok(())
1305    }
1306
1307    /// Returns the estimated total encoded bytes for this row group
1308    fn get_estimated_total_bytes(&self) -> usize {
1309        self.writers
1310            .iter()
1311            .map(|x| x.get_estimated_total_bytes())
1312            .sum()
1313    }
1314
1315    fn close(self) -> Result<Vec<ArrowColumnChunk>> {
1316        self.writers
1317            .into_iter()
1318            .map(|writer| writer.close())
1319            .collect()
1320    }
1321}
1322
1323/// Factory that creates new column writers for each row group in the Parquet file.
1324///
1325/// You can create this structure via an [`ArrowWriter::into_serialized_writer`].
1326/// See the example on [`ArrowColumnWriter`] for how to encode columns in parallel
1327#[derive(Debug)]
1328pub struct ArrowRowGroupWriterFactory {
1329    schema: SchemaDescPtr,
1330    arrow_schema: SchemaRef,
1331    props: WriterPropertiesPtr,
1332    page_store_factory: Arc<dyn PageStoreFactory>,
1333    #[cfg(feature = "encryption")]
1334    file_encryptor: Option<Arc<FileEncryptor>>,
1335}
1336
1337impl ArrowRowGroupWriterFactory {
1338    /// Create a new [`ArrowRowGroupWriterFactory`] for the provided file writer and Arrow schema
1339    pub fn new<W: Write + Send>(
1340        file_writer: &SerializedFileWriter<W>,
1341        arrow_schema: SchemaRef,
1342    ) -> Self {
1343        let schema = Arc::clone(file_writer.schema_descr_ptr());
1344        let props = Arc::clone(file_writer.properties());
1345        Self {
1346            schema,
1347            arrow_schema,
1348            props,
1349            page_store_factory: Arc::new(InMemoryPageStoreFactory),
1350            #[cfg(feature = "encryption")]
1351            file_encryptor: file_writer.file_encryptor(),
1352        }
1353    }
1354
1355    /// Set the [`PageStoreFactory`] used to allocate the buffer for each column
1356    /// chunk, e.g. to spill completed pages to a temp file or object storage
1357    /// instead of the heap. Defaults to [`InMemoryPageStoreFactory`].
1358    pub fn with_page_store_factory(
1359        mut self,
1360        page_store_factory: Arc<dyn PageStoreFactory>,
1361    ) -> Self {
1362        self.page_store_factory = page_store_factory;
1363        self
1364    }
1365
1366    fn create_row_group_writer(&self, row_group_index: usize) -> Result<ArrowRowGroupWriter> {
1367        let writers = self.create_column_writers(row_group_index)?;
1368        Ok(ArrowRowGroupWriter::new(writers, &self.arrow_schema))
1369    }
1370
1371    /// Create column writers for a new row group, with the given row group index
1372    pub fn create_column_writers(&self, row_group_index: usize) -> Result<Vec<ArrowColumnWriter>> {
1373        let mut writers = Vec::with_capacity(self.arrow_schema.fields.len());
1374        let mut leaves = self.schema.columns().iter();
1375        let column_factory = self.column_writer_factory(row_group_index);
1376        for field in &self.arrow_schema.fields {
1377            column_factory.get_arrow_column_writer(
1378                field.data_type(),
1379                &self.props,
1380                &mut leaves,
1381                &mut writers,
1382            )?;
1383        }
1384        Ok(writers)
1385    }
1386
1387    #[cfg(feature = "encryption")]
1388    fn column_writer_factory(&self, row_group_idx: usize) -> ArrowColumnWriterFactory {
1389        ArrowColumnWriterFactory::new()
1390            .with_page_store_factory(self.page_store_factory.clone())
1391            .with_file_encryptor(row_group_idx, self.file_encryptor.clone())
1392    }
1393
1394    #[cfg(not(feature = "encryption"))]
1395    fn column_writer_factory(&self, _row_group_idx: usize) -> ArrowColumnWriterFactory {
1396        ArrowColumnWriterFactory::new().with_page_store_factory(self.page_store_factory.clone())
1397    }
1398}
1399
1400/// Creates [`ArrowColumnWriter`] instances
1401struct ArrowColumnWriterFactory {
1402    /// Allocates the per-column-chunk [`PageStore`] backing each page writer.
1403    page_store_factory: Arc<dyn PageStoreFactory>,
1404    #[cfg(feature = "encryption")]
1405    row_group_index: usize,
1406    #[cfg(feature = "encryption")]
1407    file_encryptor: Option<Arc<FileEncryptor>>,
1408}
1409
1410impl ArrowColumnWriterFactory {
1411    pub fn new() -> Self {
1412        Self {
1413            page_store_factory: Arc::new(InMemoryPageStoreFactory),
1414            #[cfg(feature = "encryption")]
1415            row_group_index: 0,
1416            #[cfg(feature = "encryption")]
1417            file_encryptor: None,
1418        }
1419    }
1420
1421    /// Use `page_store_factory` to allocate the buffer for each column chunk.
1422    pub fn with_page_store_factory(
1423        mut self,
1424        page_store_factory: Arc<dyn PageStoreFactory>,
1425    ) -> Self {
1426        self.page_store_factory = page_store_factory;
1427        self
1428    }
1429
1430    #[cfg(feature = "encryption")]
1431    pub fn with_file_encryptor(
1432        mut self,
1433        row_group_index: usize,
1434        file_encryptor: Option<Arc<FileEncryptor>>,
1435    ) -> Self {
1436        self.row_group_index = row_group_index;
1437        self.file_encryptor = file_encryptor;
1438        self
1439    }
1440
1441    #[cfg(feature = "encryption")]
1442    fn create_page_writer(
1443        &self,
1444        column_descriptor: &ColumnDescPtr,
1445        column_index: usize,
1446    ) -> Result<Box<ArrowPageWriter>> {
1447        let column_path = column_descriptor.path().string();
1448        let page_encryptor = PageEncryptor::create_if_column_encrypted(
1449            self.file_encryptor.as_ref(),
1450            self.row_group_index,
1451            column_index,
1452            &column_path,
1453        )?;
1454        let args = PageStoreArgs::new(column_index, column_descriptor);
1455        let store = self.page_store_factory.create(&args)?;
1456        Ok(Box::new(
1457            ArrowPageWriter::new(store).with_encryptor(page_encryptor),
1458        ))
1459    }
1460
1461    #[cfg(not(feature = "encryption"))]
1462    fn create_page_writer(
1463        &self,
1464        column_descriptor: &ColumnDescPtr,
1465        column_index: usize,
1466    ) -> Result<Box<ArrowPageWriter>> {
1467        let args = PageStoreArgs::new(column_index, column_descriptor);
1468        let store = self.page_store_factory.create(&args)?;
1469        Ok(Box::new(ArrowPageWriter::new(store)))
1470    }
1471
1472    /// Gets an [`ArrowColumnWriter`] for the given `data_type`, appending the
1473    /// output ColumnDesc to `leaves` and the column writers to `out`
1474    fn get_arrow_column_writer(
1475        &self,
1476        data_type: &ArrowDataType,
1477        props: &WriterPropertiesPtr,
1478        leaves: &mut Iter<'_, ColumnDescPtr>,
1479        out: &mut Vec<ArrowColumnWriter>,
1480    ) -> Result<()> {
1481        let write_distinct_values = props.write_row_group_number_distinct_values();
1482
1483        // Instantiate writers for normal columns
1484        let col = |desc: &ColumnDescPtr| -> Result<ArrowColumnWriter> {
1485            let page_writer = self.create_page_writer(desc, out.len())?;
1486            let chunk = page_writer.buffer.clone();
1487            let writer = get_column_writer(desc.clone(), props.clone(), page_writer);
1488            Ok(ArrowColumnWriter {
1489                chunk,
1490                writer: ArrowColumnWriterImpl::Column(writer),
1491                distinct_values_seen: write_distinct_values.then(HashSet::new),
1492            })
1493        };
1494
1495        // Instantiate writers for byte arrays (e.g. Utf8,  Binary, etc)
1496        let bytes = |desc: &ColumnDescPtr| -> Result<ArrowColumnWriter> {
1497            let page_writer = self.create_page_writer(desc, out.len())?;
1498            let chunk = page_writer.buffer.clone();
1499            let writer = GenericColumnWriter::new(desc.clone(), props.clone(), page_writer);
1500            Ok(ArrowColumnWriter {
1501                chunk,
1502                writer: ArrowColumnWriterImpl::ByteArray(writer),
1503                distinct_values_seen: write_distinct_values.then(HashSet::new),
1504            })
1505        };
1506
1507        match data_type {
1508            _ if data_type.is_primitive() => out.push(col(leaves.next().unwrap())?),
1509            ArrowDataType::FixedSizeBinary(_) | ArrowDataType::Boolean | ArrowDataType::Null => {
1510                out.push(col(leaves.next().unwrap())?)
1511            }
1512            ArrowDataType::LargeBinary
1513            | ArrowDataType::Binary
1514            | ArrowDataType::Utf8
1515            | ArrowDataType::LargeUtf8
1516            | ArrowDataType::BinaryView
1517            | ArrowDataType::Utf8View => out.push(bytes(leaves.next().unwrap())?),
1518            ArrowDataType::List(f)
1519            | ArrowDataType::LargeList(f)
1520            | ArrowDataType::FixedSizeList(f, _)
1521            | ArrowDataType::ListView(f)
1522            | ArrowDataType::LargeListView(f) => {
1523                self.get_arrow_column_writer(f.data_type(), props, leaves, out)?
1524            }
1525            ArrowDataType::Struct(fields) => {
1526                for field in fields {
1527                    self.get_arrow_column_writer(field.data_type(), props, leaves, out)?
1528                }
1529            }
1530            ArrowDataType::Map(f, _) => match f.data_type() {
1531                ArrowDataType::Struct(f) => {
1532                    self.get_arrow_column_writer(f[0].data_type(), props, leaves, out)?;
1533                    self.get_arrow_column_writer(f[1].data_type(), props, leaves, out)?
1534                }
1535                _ => unreachable!("invalid map type"),
1536            },
1537            ArrowDataType::Dictionary(_, value_type) => match value_type.as_ref() {
1538                ArrowDataType::Utf8
1539                | ArrowDataType::LargeUtf8
1540                | ArrowDataType::Binary
1541                | ArrowDataType::LargeBinary => out.push(bytes(leaves.next().unwrap())?),
1542                ArrowDataType::Utf8View | ArrowDataType::BinaryView => {
1543                    out.push(bytes(leaves.next().unwrap())?)
1544                }
1545                ArrowDataType::FixedSizeBinary(_) => out.push(bytes(leaves.next().unwrap())?),
1546                _ => out.push(col(leaves.next().unwrap())?),
1547            },
1548            ArrowDataType::RunEndEncoded(_, value_field) => {
1549                self.get_arrow_column_writer(value_field.data_type(), props, leaves, out)?
1550            }
1551            _ => {
1552                return Err(ParquetError::NYI(format!(
1553                    "Attempting to write an Arrow type {data_type} to parquet that is not yet implemented"
1554                )));
1555            }
1556        }
1557        Ok(())
1558    }
1559}
1560
1561fn write_leaf(
1562    writer: &mut ColumnWriter<'_>,
1563    column: &dyn arrow_array::Array,
1564    levels: &ArrayLevels,
1565) -> Result<usize> {
1566    let indices = levels.non_null_indices();
1567
1568    match writer {
1569        // Note: this should match the contents of arrow_to_parquet_type
1570        ColumnWriter::Int32ColumnWriter(typed) => {
1571            match column.data_type() {
1572                ArrowDataType::Null => {
1573                    let array = Int32Array::new_null(column.len());
1574                    write_primitive(typed, array.values(), levels)
1575                }
1576                ArrowDataType::Int8 => {
1577                    let array: Int32Array = column.as_primitive::<Int8Type>().unary(|x| x as i32);
1578                    write_primitive(typed, array.values(), levels)
1579                }
1580                ArrowDataType::Int16 => {
1581                    let array: Int32Array = column.as_primitive::<Int16Type>().unary(|x| x as i32);
1582                    write_primitive(typed, array.values(), levels)
1583                }
1584                ArrowDataType::Int32 => {
1585                    write_primitive(typed, column.as_primitive::<Int32Type>().values(), levels)
1586                }
1587                ArrowDataType::UInt8 => {
1588                    let array: Int32Array = column.as_primitive::<UInt8Type>().unary(|x| x as i32);
1589                    write_primitive(typed, array.values(), levels)
1590                }
1591                ArrowDataType::UInt16 => {
1592                    let array: Int32Array = column.as_primitive::<UInt16Type>().unary(|x| x as i32);
1593                    write_primitive(typed, array.values(), levels)
1594                }
1595                ArrowDataType::UInt32 => {
1596                    // follow C++ implementation and use overflow/reinterpret cast from  u32 to i32 which will map
1597                    // `(i32::MAX as u32)..u32::MAX` to `i32::MIN..0`
1598                    let array = column.as_primitive::<UInt32Type>();
1599                    write_primitive(typed, array.values().inner().typed_data(), levels)
1600                }
1601                ArrowDataType::Date32 => {
1602                    let array = column.as_primitive::<Date32Type>();
1603                    write_primitive(typed, array.values(), levels)
1604                }
1605                ArrowDataType::Time32(TimeUnit::Second) => {
1606                    let array = column.as_primitive::<Time32SecondType>();
1607                    write_primitive(typed, array.values(), levels)
1608                }
1609                ArrowDataType::Time32(TimeUnit::Millisecond) => {
1610                    let array = column.as_primitive::<Time32MillisecondType>();
1611                    write_primitive(typed, array.values(), levels)
1612                }
1613                ArrowDataType::Date64 => {
1614                    // If the column is a Date64, we truncate it
1615                    let array: Int32Array = column
1616                        .as_primitive::<Date64Type>()
1617                        .unary(|x| (x / 86_400_000) as _);
1618
1619                    write_primitive(typed, array.values(), levels)
1620                }
1621                ArrowDataType::Decimal32(_, _) => {
1622                    let array = column
1623                        .as_primitive::<Decimal32Type>()
1624                        .unary::<_, Int32Type>(|v| v);
1625                    write_primitive(typed, array.values(), levels)
1626                }
1627                ArrowDataType::Decimal64(_, _) => {
1628                    // use the int32 to represent the decimal with low precision
1629                    let array = column
1630                        .as_primitive::<Decimal64Type>()
1631                        .unary::<_, Int32Type>(|v| v as i32);
1632                    write_primitive(typed, array.values(), levels)
1633                }
1634                ArrowDataType::Decimal128(_, _) => {
1635                    // use the int32 to represent the decimal with low precision
1636                    let array = column
1637                        .as_primitive::<Decimal128Type>()
1638                        .unary::<_, Int32Type>(|v| v as i32);
1639                    write_primitive(typed, array.values(), levels)
1640                }
1641                ArrowDataType::Decimal256(_, _) => {
1642                    // use the int32 to represent the decimal with low precision
1643                    let array = column
1644                        .as_primitive::<Decimal256Type>()
1645                        .unary::<_, Int32Type>(|v| v.as_i128() as i32);
1646                    write_primitive(typed, array.values(), levels)
1647                }
1648                d => Err(ParquetError::General(format!("Cannot coerce {d} to I32"))),
1649            }
1650        }
1651        ColumnWriter::BoolColumnWriter(typed) => {
1652            let array = column.as_boolean();
1653            let values = get_bool_array_slice(array, indices.iter().copied());
1654            typed.write_batch_internal(
1655                values.as_slice(),
1656                None,
1657                levels.def_level_data().as_ref(),
1658                levels.rep_level_data().as_ref(),
1659                None,
1660                None,
1661                None,
1662            )
1663        }
1664        ColumnWriter::Int64ColumnWriter(typed) => {
1665            match column.data_type() {
1666                ArrowDataType::Date64 => {
1667                    let array = column
1668                        .as_primitive::<Date64Type>()
1669                        .reinterpret_cast::<Int64Type>();
1670
1671                    write_primitive(typed, array.values(), levels)
1672                }
1673                ArrowDataType::Int64 => {
1674                    let array = column.as_primitive::<Int64Type>();
1675                    write_primitive(typed, array.values(), levels)
1676                }
1677                ArrowDataType::UInt64 => {
1678                    let values = column.as_primitive::<UInt64Type>().values();
1679                    // follow C++ implementation and use overflow/reinterpret cast from  u64 to i64 which will map
1680                    // `(i64::MAX as u64)..u64::MAX` to `i64::MIN..0`
1681                    let array = values.inner().typed_data::<i64>();
1682                    write_primitive(typed, array, levels)
1683                }
1684                ArrowDataType::Time64(TimeUnit::Microsecond) => {
1685                    let array = column.as_primitive::<Time64MicrosecondType>();
1686                    write_primitive(typed, array.values(), levels)
1687                }
1688                ArrowDataType::Time64(TimeUnit::Nanosecond) => {
1689                    let array = column.as_primitive::<Time64NanosecondType>();
1690                    write_primitive(typed, array.values(), levels)
1691                }
1692                ArrowDataType::Timestamp(unit, _) => match unit {
1693                    TimeUnit::Second => {
1694                        let array = column.as_primitive::<TimestampSecondType>();
1695                        write_primitive(typed, array.values(), levels)
1696                    }
1697                    TimeUnit::Millisecond => {
1698                        let array = column.as_primitive::<TimestampMillisecondType>();
1699                        write_primitive(typed, array.values(), levels)
1700                    }
1701                    TimeUnit::Microsecond => {
1702                        let array = column.as_primitive::<TimestampMicrosecondType>();
1703                        write_primitive(typed, array.values(), levels)
1704                    }
1705                    TimeUnit::Nanosecond => {
1706                        let array = column.as_primitive::<TimestampNanosecondType>();
1707                        write_primitive(typed, array.values(), levels)
1708                    }
1709                },
1710                ArrowDataType::Duration(unit) => match unit {
1711                    TimeUnit::Second => {
1712                        let array = column.as_primitive::<DurationSecondType>();
1713                        write_primitive(typed, array.values(), levels)
1714                    }
1715                    TimeUnit::Millisecond => {
1716                        let array = column.as_primitive::<DurationMillisecondType>();
1717                        write_primitive(typed, array.values(), levels)
1718                    }
1719                    TimeUnit::Microsecond => {
1720                        let array = column.as_primitive::<DurationMicrosecondType>();
1721                        write_primitive(typed, array.values(), levels)
1722                    }
1723                    TimeUnit::Nanosecond => {
1724                        let array = column.as_primitive::<DurationNanosecondType>();
1725                        write_primitive(typed, array.values(), levels)
1726                    }
1727                },
1728                ArrowDataType::Decimal64(_, _) => {
1729                    let array = column
1730                        .as_primitive::<Decimal64Type>()
1731                        .reinterpret_cast::<Int64Type>();
1732                    write_primitive(typed, array.values(), levels)
1733                }
1734                ArrowDataType::Decimal128(_, _) => {
1735                    // use the int64 to represent the decimal with low precision
1736                    let array = column
1737                        .as_primitive::<Decimal128Type>()
1738                        .unary::<_, Int64Type>(|v| v as i64);
1739                    write_primitive(typed, array.values(), levels)
1740                }
1741                ArrowDataType::Decimal256(_, _) => {
1742                    // use the int64 to represent the decimal with low precision
1743                    let array = column
1744                        .as_primitive::<Decimal256Type>()
1745                        .unary::<_, Int64Type>(|v| v.as_i128() as i64);
1746                    write_primitive(typed, array.values(), levels)
1747                }
1748                d => Err(ParquetError::General(format!("Cannot coerce {d} to I64"))),
1749            }
1750        }
1751        ColumnWriter::Int96ColumnWriter(_typed) => {
1752            unreachable!("Currently unreachable because data type not supported")
1753        }
1754        ColumnWriter::FloatColumnWriter(typed) => {
1755            let array = column.as_primitive::<Float32Type>();
1756            write_primitive(typed, array.values(), levels)
1757        }
1758        ColumnWriter::DoubleColumnWriter(typed) => {
1759            let array = column.as_primitive::<Float64Type>();
1760            write_primitive(typed, array.values(), levels)
1761        }
1762        ColumnWriter::ByteArrayColumnWriter(_) => {
1763            unreachable!("should use ByteArrayWriter")
1764        }
1765        ColumnWriter::FixedLenByteArrayColumnWriter(typed) => {
1766            let bytes = match column.data_type() {
1767                ArrowDataType::Interval(interval_unit) => match interval_unit {
1768                    IntervalUnit::YearMonth => {
1769                        let array = column.as_primitive::<IntervalYearMonthType>();
1770                        get_interval_ym_array_slice(array, indices.iter().copied())
1771                    }
1772                    IntervalUnit::DayTime => {
1773                        let array = column.as_primitive::<IntervalDayTimeType>();
1774                        get_interval_dt_array_slice(array, indices.iter().copied())
1775                    }
1776                    IntervalUnit::MonthDayNano => {
1777                        return Err(ParquetError::NYI(format!(
1778                            "Attempting to write an Arrow interval type {interval_unit:?} to parquet that is not yet implemented"
1779                        )));
1780                    }
1781                },
1782                ArrowDataType::FixedSizeBinary(_) => {
1783                    let array = column.as_fixed_size_binary();
1784                    get_fsb_array_slice(array, indices.iter().copied())
1785                }
1786                ArrowDataType::Decimal32(_, _) => {
1787                    let array = column.as_primitive::<Decimal32Type>();
1788                    get_decimal_array_slice(array, indices.iter().copied())
1789                }
1790                ArrowDataType::Decimal64(_, _) => {
1791                    let array = column.as_primitive::<Decimal64Type>();
1792                    get_decimal_array_slice(array, indices.iter().copied())
1793                }
1794                ArrowDataType::Decimal128(_, _) => {
1795                    let array = column.as_primitive::<Decimal128Type>();
1796                    get_decimal_array_slice(array, indices.iter().copied())
1797                }
1798                ArrowDataType::Decimal256(_, _) => {
1799                    let array = column.as_primitive::<Decimal256Type>();
1800                    get_decimal_array_slice(array, indices.iter().copied())
1801                }
1802                ArrowDataType::Float16 => {
1803                    let array = column.as_primitive::<Float16Type>();
1804                    get_float_16_array_slice(array, indices.iter().copied())
1805                }
1806                _ => {
1807                    return Err(ParquetError::NYI(
1808                        "Attempting to write an Arrow type that is not yet implemented".to_string(),
1809                    ));
1810                }
1811            };
1812            typed.write_batch_internal(
1813                bytes.as_slice(),
1814                None,
1815                levels.def_level_data().as_ref(),
1816                levels.rep_level_data().as_ref(),
1817                None,
1818                None,
1819                None,
1820            )
1821        }
1822    }
1823}
1824
1825fn write_primitive<E: ColumnValueEncoder>(
1826    writer: &mut GenericColumnWriter<E>,
1827    values: &E::Values,
1828    levels: &ArrayLevels,
1829) -> Result<usize> {
1830    writer.write_batch_internal(
1831        values,
1832        Some(levels.non_null_indices()),
1833        levels.def_level_data().as_ref(),
1834        levels.rep_level_data().as_ref(),
1835        None,
1836        None,
1837        None,
1838    )
1839}
1840
1841fn get_bool_array_slice(
1842    array: &arrow_array::BooleanArray,
1843    indices: impl ExactSizeIterator<Item = usize>,
1844) -> Vec<bool> {
1845    let mut values = Vec::with_capacity(indices.len());
1846    for i in indices {
1847        values.push(array.value(i))
1848    }
1849    values
1850}
1851
1852/// Returns 12-byte values representing 3 values of months, days and milliseconds (4-bytes each).
1853/// An Arrow YearMonth interval only stores months, thus only the first 4 bytes are populated.
1854fn get_interval_ym_array_slice(
1855    array: &arrow_array::IntervalYearMonthArray,
1856    indices: impl ExactSizeIterator<Item = usize>,
1857) -> Vec<FixedLenByteArray> {
1858    chunk_array_slice(12, indices, move |i, chunk| {
1859        let value = array.value(i);
1860        chunk[0..4].copy_from_slice(&value.to_le_bytes());
1861    })
1862}
1863
1864/// Returns 12-byte values representing 3 values of months, days and milliseconds (4-bytes each).
1865/// An Arrow DayTime interval only stores days and millis, thus the first 4 bytes are not populated.
1866fn get_interval_dt_array_slice(
1867    array: &arrow_array::IntervalDayTimeArray,
1868    indices: impl ExactSizeIterator<Item = usize>,
1869) -> Vec<FixedLenByteArray> {
1870    chunk_array_slice(12, indices, move |i, chunk| {
1871        let value = array.value(i);
1872        chunk[4..8].copy_from_slice(&value.days.to_le_bytes());
1873        chunk[8..12].copy_from_slice(&value.milliseconds.to_le_bytes());
1874    })
1875}
1876
1877trait NativeDecimalType: DecimalType {
1878    type NativeBytes: AsRef<[u8]>;
1879
1880    fn to_be_bytes(value: Self::Native) -> Self::NativeBytes;
1881}
1882impl NativeDecimalType for Decimal32Type {
1883    type NativeBytes = [u8; Self::BYTE_LENGTH];
1884
1885    fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
1886        value.to_be_bytes()
1887    }
1888}
1889impl NativeDecimalType for Decimal64Type {
1890    type NativeBytes = [u8; Self::BYTE_LENGTH];
1891
1892    fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
1893        value.to_be_bytes()
1894    }
1895}
1896impl NativeDecimalType for Decimal128Type {
1897    type NativeBytes = [u8; Self::BYTE_LENGTH];
1898
1899    fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
1900        value.to_be_bytes()
1901    }
1902}
1903impl NativeDecimalType for Decimal256Type {
1904    type NativeBytes = [u8; Self::BYTE_LENGTH];
1905
1906    fn to_be_bytes(value: Self::Native) -> Self::NativeBytes {
1907        value.to_be_bytes()
1908    }
1909}
1910
1911fn get_decimal_array_slice<T: NativeDecimalType>(
1912    array: &PrimitiveArray<T>,
1913    indices: impl ExactSizeIterator<Item = usize>,
1914) -> Vec<FixedLenByteArray> {
1915    let chunk_size = decimal_length_from_precision(array.precision());
1916    assert!(chunk_size <= T::BYTE_LENGTH);
1917
1918    if chunk_size == T::BYTE_LENGTH {
1919        // Special-case that allows inlining memcpy.
1920        chunk_array_slice(chunk_size, indices, move |i, chunk| {
1921            let as_be_bytes = T::to_be_bytes(array.value(i));
1922            chunk.copy_from_slice(as_be_bytes.as_ref());
1923        })
1924    } else {
1925        chunk_array_slice(chunk_size, indices, move |i, chunk| {
1926            let as_be_bytes = T::to_be_bytes(array.value(i));
1927            let resized_value = &as_be_bytes.as_ref()[(T::BYTE_LENGTH - chunk.len())..];
1928            chunk.copy_from_slice(resized_value);
1929        })
1930    }
1931}
1932
1933fn get_float_16_array_slice(
1934    array: &arrow_array::Float16Array,
1935    indices: impl ExactSizeIterator<Item = usize>,
1936) -> Vec<FixedLenByteArray> {
1937    chunk_array_slice(2, indices, move |i, chunk| {
1938        let value = array.value(i).to_le_bytes();
1939        chunk.copy_from_slice(&value);
1940    })
1941}
1942
1943fn get_fsb_array_slice(
1944    array: &arrow_array::FixedSizeBinaryArray,
1945    indices: impl ExactSizeIterator<Item = usize>,
1946) -> Vec<FixedLenByteArray> {
1947    chunk_array_slice(array.value_size(), indices, move |i, chunk| {
1948        let value = array.value(i);
1949        chunk.copy_from_slice(value);
1950    })
1951}
1952
1953#[inline]
1954fn chunk_array_slice(
1955    chunk_size: usize,
1956    indices: impl ExactSizeIterator<Item = usize>,
1957    writer: impl Fn(usize, &mut [u8]),
1958) -> Vec<FixedLenByteArray> {
1959    let capacity = indices.len() * chunk_size;
1960    // TODO: This could be done with Vec::spare_capacity_mut,
1961    //       but [MaybeUninit]::write_copy_of_slice is gated behind MSRV 1.93
1962    let mut arena = vec![0; capacity];
1963    for (i, chunk) in indices.zip(arena.chunks_exact_mut(chunk_size)) {
1964        writer(i, chunk);
1965    }
1966    chunk_contiguous_vec(arena, chunk_size)
1967}
1968
1969fn chunk_contiguous_vec(arena: Vec<u8>, chunk_size: usize) -> Vec<FixedLenByteArray> {
1970    let mut values = Vec::with_capacity(arena.len() / chunk_size);
1971    let mut arena = Bytes::from(arena);
1972    while arena.len() >= chunk_size {
1973        let slice = arena.split_to(chunk_size);
1974        values.push(FixedLenByteArray::from(ByteArray::from(slice)));
1975    }
1976    values
1977}
1978
1979/// Hash a byte slice to a u64 for NDV tracking.
1980#[inline]
1981fn hash_bytes(bytes: &[u8]) -> u64 {
1982    twox_hash::XxHash64::oneshot(0, bytes)
1983}
1984
1985/// Returns the byte width of an Arrow dictionary key type, or 0 if unsupported.
1986fn arrow_key_byte_width(dt: &ArrowDataType) -> usize {
1987    match dt {
1988        ArrowDataType::Int8 | ArrowDataType::UInt8 => 1,
1989        ArrowDataType::Int16 | ArrowDataType::UInt16 => 2,
1990        ArrowDataType::Int32 | ArrowDataType::UInt32 => 4,
1991        ArrowDataType::Int64 | ArrowDataType::UInt64 => 8,
1992        _ => 0,
1993    }
1994}
1995
1996/// Returns the fixed byte width for primitive Arrow types, or `None` for variable-length types.
1997fn fixed_byte_width(dt: &ArrowDataType) -> Option<usize> {
1998    use ArrowDataType::*;
1999    match dt {
2000        Int8 | UInt8 => Some(1),
2001        Int16 | UInt16 | Float16 => Some(2),
2002        Int32 | UInt32 | Float32 | Date32 | Time32(_) | Decimal32(_, _) => Some(4),
2003        Int64
2004        | UInt64
2005        | Float64
2006        | Date64
2007        | Time64(_)
2008        | Timestamp(_, _)
2009        | Duration(_)
2010        | Decimal64(_, _) => Some(8),
2011        Interval(IntervalUnit::YearMonth) => Some(4),
2012        Interval(IntervalUnit::DayTime) => Some(8),
2013        Interval(IntervalUnit::MonthDayNano) => Some(16),
2014        Decimal128(_, _) => Some(16),
2015        Decimal256(_, _) => Some(32),
2016        _ => None,
2017    }
2018}
2019
2020/// Hash the non-null values in `array` (at `non_null_indices`) into `seen`.
2021///
2022/// Handles primitive, boolean, fixed-size-binary, and variable-length (Utf8/Binary)
2023/// arrays. Unsupported types are silently skipped, leaving `seen` unchanged for
2024/// those values (NDV is best-effort).
2025fn update_distinct_values_seen(
2026    array: &dyn arrow_array::Array,
2027    non_null_indices: &[usize],
2028    seen: &mut DistinctValuesSet,
2029) {
2030    let data = array.to_data();
2031    let offset = data.offset();
2032
2033    match array.data_type() {
2034        ArrowDataType::Boolean => {
2035            let arr = array
2036                .as_any()
2037                .downcast_ref::<arrow_array::BooleanArray>()
2038                .unwrap();
2039            for &row in non_null_indices {
2040                seen.insert(arr.value(row) as u64);
2041            }
2042        }
2043        ArrowDataType::Utf8 | ArrowDataType::Binary => {
2044            let offsets = data.buffers()[0].typed_data::<i32>();
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::LargeUtf8 | ArrowDataType::LargeBinary => {
2053            let offsets = data.buffers()[0].typed_data::<i64>();
2054            let values = data.buffers()[1].as_slice();
2055            for &row in non_null_indices {
2056                let start = offsets[offset + row] as usize;
2057                let end = offsets[offset + row + 1] as usize;
2058                seen.insert(hash_bytes(&values[start..end]));
2059            }
2060        }
2061        ArrowDataType::FixedSizeBinary(byte_width) => {
2062            let byte_width = *byte_width as usize;
2063            let buffer = data.buffers()[0].as_slice();
2064            for &row in non_null_indices {
2065                let start = (offset + row) * byte_width;
2066                seen.insert(hash_bytes(&buffer[start..start + byte_width]));
2067            }
2068        }
2069        ArrowDataType::Utf8View => {
2070            let string_view_array = array.as_string_view();
2071            for &row in non_null_indices {
2072                seen.insert(hash_bytes(string_view_array.value(row).as_bytes()));
2073            }
2074        }
2075        ArrowDataType::BinaryView => {
2076            let binary_view_array = array.as_binary_view();
2077            for &row in non_null_indices {
2078                seen.insert(hash_bytes(binary_view_array.value(row)));
2079            }
2080        }
2081        data_type => {
2082            if let Some(width) = fixed_byte_width(data_type) {
2083                let buffer = data.buffers()[0].as_slice();
2084                for &row in non_null_indices {
2085                    let pos = (offset + row) * width;
2086                    seen.insert(hash_bytes(&buffer[pos..pos + width]));
2087                }
2088            }
2089            // nested types (List, LargeList, etc.) are Parquet groups, not leaf columns: skip
2090        }
2091    }
2092}
2093
2094#[cfg(test)]
2095mod tests {
2096    use super::*;
2097    use std::cmp::Ordering;
2098    use std::collections::HashMap;
2099
2100    use std::fs::File;
2101
2102    use crate::arrow::arrow_reader::{ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder};
2103    use crate::arrow::{ARROW_SCHEMA_META_KEY, PARQUET_FIELD_ID_META_KEY};
2104    use crate::column::page::{Page, PageReader};
2105    use crate::file::metadata::thrift::PageHeader;
2106    use crate::file::page_index::column_index::ColumnIndexMetaData;
2107    use crate::file::reader::SerializedPageReader;
2108    use crate::parquet_thrift::{ReadThrift, ThriftSliceInputProtocol};
2109    use crate::schema::types::ColumnPath;
2110    use arrow::datatypes::ToByteSlice;
2111    use arrow::datatypes::{DataType, Schema};
2112    use arrow::error::Result as ArrowResult;
2113    use arrow::util::data_gen::create_random_array;
2114    use arrow::util::pretty::pretty_format_batches;
2115    use arrow::{array::*, buffer::Buffer};
2116    use arrow_buffer::{IntervalDayTime, IntervalMonthDayNano, NullBuffer, OffsetBuffer, i256};
2117    use arrow_schema::Fields;
2118    use half::f16;
2119    use num_traits::{FromPrimitive, ToPrimitive};
2120    use tempfile::tempfile;
2121
2122    use crate::basic::{Encoding, EncodingMask};
2123    use crate::data_type::AsBytes;
2124    use crate::file::metadata::{ColumnChunkMetaData, ParquetMetaData, ParquetMetaDataReader};
2125    use crate::file::properties::{
2126        BloomFilterPosition, EnabledStatistics, ReaderProperties, WriterVersion,
2127    };
2128    use crate::file::serialized_reader::ReadOptionsBuilder;
2129    use crate::file::{
2130        reader::{FileReader, SerializedFileReader},
2131        statistics::Statistics,
2132    };
2133
2134    /// A [`PageStore`] that allocates *sparse, non-contiguous* handles and keeps
2135    /// blobs in a `HashMap` — nothing like the default `Vec<Bytes>`. Used to
2136    /// prove the writer relies only on the opaque-handle contract and never on
2137    /// handles being dense `Vec` indices. Records how many blobs were stored.
2138    #[derive(Debug, Default)]
2139    struct RecordingPageStore {
2140        next: u64,
2141        blobs: HashMap<u64, Bytes>,
2142        puts: Arc<std::sync::atomic::AtomicUsize>,
2143    }
2144
2145    impl PageStore for RecordingPageStore {
2146        fn put(&mut self, value: Bytes) -> Result<PageKey> {
2147            // Deliberately non-sequential, never-zero handles.
2148            let id = 100 + self.next * 7;
2149            self.next += 1;
2150            self.puts.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2151            self.blobs.insert(id, value);
2152            Ok(PageKey::new(id))
2153        }
2154
2155        fn take(&mut self, key: PageKey) -> Result<Bytes> {
2156            self.blobs
2157                .remove(&key.get())
2158                .ok_or_else(|| ParquetError::General(format!("missing key {}", key.get())))
2159        }
2160    }
2161
2162    #[derive(Debug)]
2163    struct RecordingPageStoreFactory {
2164        puts: Arc<std::sync::atomic::AtomicUsize>,
2165    }
2166
2167    impl PageStoreFactory for RecordingPageStoreFactory {
2168        fn create(&self, _args: &PageStoreArgs<'_>) -> Result<Box<dyn PageStore>> {
2169            Ok(Box::new(RecordingPageStore {
2170                puts: self.puts.clone(),
2171                ..Default::default()
2172            }))
2173        }
2174    }
2175
2176    /// A custom [`PageStore`] must produce byte-identical files to the in-memory
2177    /// default, across dictionary and non-dictionary columns and multiple row
2178    /// groups (so multiple store instances are exercised).
2179    #[test]
2180    fn custom_page_store_is_byte_identical_to_default() {
2181        let schema = Arc::new(Schema::new(vec![
2182            Field::new("i", DataType::Int32, true),
2183            // A low-cardinality string column to exercise the dictionary path.
2184            Field::new("s", DataType::Utf8, true),
2185        ]));
2186        let i = Int32Array::from(vec![Some(1), None, Some(3), Some(4), Some(5), Some(6)]);
2187        let s = StringArray::from(vec![
2188            Some("a"),
2189            Some("bb"),
2190            Some("a"),
2191            None,
2192            Some("bb"),
2193            Some("ccc"),
2194        ]);
2195        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(i), Arc::new(s)]).unwrap();
2196
2197        // Small row groups so multiple column chunks (hence multiple store
2198        // instances) are produced.
2199        let props = WriterProperties::builder()
2200            .set_max_row_group_row_count(Some(3))
2201            .build();
2202
2203        let write = |factory: Option<Arc<dyn PageStoreFactory>>| {
2204            let mut buffer = Vec::new();
2205            let mut opts = ArrowWriterOptions::new().with_properties(props.clone());
2206            if let Some(factory) = factory {
2207                opts = opts.with_page_store_factory(factory);
2208            }
2209            let mut writer =
2210                ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap();
2211            writer.write(&batch).unwrap();
2212            writer.close().unwrap();
2213            buffer
2214        };
2215
2216        let default_bytes = write(None);
2217
2218        let puts = Arc::new(std::sync::atomic::AtomicUsize::new(0));
2219        let custom_bytes = write(Some(Arc::new(RecordingPageStoreFactory {
2220            puts: puts.clone(),
2221        })));
2222
2223        assert!(
2224            puts.load(std::sync::atomic::Ordering::Relaxed) > 0,
2225            "custom PageStore was never written to"
2226        );
2227        assert_eq!(
2228            default_bytes, custom_bytes,
2229            "a custom PageStore must produce byte-identical output to the default"
2230        );
2231    }
2232
2233    /// A dictionary-encoded column written through the deferred-ordering Arrow
2234    /// path must round-trip correctly even with the offset index disabled, when
2235    /// only the chunk-level dictionary/data page offsets are rewritten (there is
2236    /// no offset index to rebuild). Spans multiple data pages so the
2237    /// dictionary-first reordering is exercised.
2238    #[test]
2239    #[cfg_attr(miri, ignore)] // Takes too long
2240    fn dictionary_column_round_trips_with_offset_index_disabled() {
2241        let schema = Arc::new(Schema::new(vec![Field::new("k", DataType::Int32, true)]));
2242
2243        // Low cardinality so the column stays dictionary-encoded; enough rows to
2244        // span several data pages within a single row group.
2245        let values: Vec<Option<i32>> = (0..50_000).map(|i| Some(i % 8)).collect();
2246        let array = Int32Array::from(values.clone());
2247        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
2248
2249        let props = WriterProperties::builder()
2250            .set_offset_index_disabled(true)
2251            .set_data_page_row_count_limit(4096)
2252            .build();
2253        let opts = ArrowWriterOptions::new().with_properties(props);
2254
2255        let mut buffer = Vec::new();
2256        let mut writer =
2257            ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap();
2258        writer.write(&batch).unwrap();
2259        writer.close().unwrap();
2260
2261        let reader = ParquetRecordBatchReader::try_new(Bytes::from(buffer), values.len()).unwrap();
2262        let read: Vec<RecordBatch> = reader.collect::<ArrowResult<_>>().unwrap();
2263        let read_values: Vec<Option<i32>> = read
2264            .iter()
2265            .flat_map(|b| b.column(0).as_primitive::<Int32Type>().iter())
2266            .collect();
2267        assert_eq!(read_values, values);
2268    }
2269
2270    /// The dictionary page is routed through the [`PageStore`] like any other
2271    /// page rather than held resident in memory, so a dictionary column chunk's
2272    /// *entire* serialized size — dictionary page included — passes through the
2273    /// store.
2274    #[test]
2275    fn dictionary_page_is_routed_through_the_store() {
2276        /// A store that sums the bytes handed to `put`.
2277        #[derive(Debug, Default)]
2278        struct SizeRecordingPageStore {
2279            blobs: Vec<Bytes>,
2280            bytes_put: Arc<std::sync::atomic::AtomicUsize>,
2281        }
2282        impl PageStore for SizeRecordingPageStore {
2283            fn put(&mut self, value: Bytes) -> Result<PageKey> {
2284                self.bytes_put
2285                    .fetch_add(value.len(), std::sync::atomic::Ordering::Relaxed);
2286                let key = PageKey::new(self.blobs.len() as u64);
2287                self.blobs.push(value);
2288                Ok(key)
2289            }
2290            fn take(&mut self, key: PageKey) -> Result<Bytes> {
2291                Ok(std::mem::take(&mut self.blobs[key.get() as usize]))
2292            }
2293        }
2294        #[derive(Debug)]
2295        struct Factory {
2296            bytes_put: Arc<std::sync::atomic::AtomicUsize>,
2297        }
2298        impl PageStoreFactory for Factory {
2299            fn create(&self, _args: &PageStoreArgs<'_>) -> Result<Box<dyn PageStore>> {
2300                Ok(Box::new(SizeRecordingPageStore {
2301                    bytes_put: self.bytes_put.clone(),
2302                    ..Default::default()
2303                }))
2304            }
2305        }
2306
2307        let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8, false)]));
2308        // Low cardinality keeps the column dictionary-encoded with a real,
2309        // non-empty dictionary page.
2310        let values: Vec<&str> = (0..2048)
2311            .map(|i| ["alpha", "beta", "gamma", "delta"][i % 4])
2312            .collect();
2313        let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(StringArray::from(values))])
2314            .unwrap();
2315
2316        let bytes_put = Arc::new(std::sync::atomic::AtomicUsize::new(0));
2317        let opts = ArrowWriterOptions::new().with_page_store_factory(Arc::new(Factory {
2318            bytes_put: bytes_put.clone(),
2319        }));
2320
2321        // A single batch / single column means exactly one row group and one
2322        // store instance, so the bytes it saw map to one column chunk.
2323        let mut buffer = Vec::new();
2324        let mut writer =
2325            ArrowWriter::try_new_with_options(&mut buffer, schema.clone(), opts).unwrap();
2326        writer.write(&batch).unwrap();
2327        writer.close().unwrap();
2328
2329        let reader = SerializedFileReader::new(Bytes::from(buffer)).unwrap();
2330        let column = reader.metadata().row_group(0).column(0);
2331        assert!(
2332            column.dictionary_page_offset().is_some(),
2333            "expected the column to be dictionary-encoded"
2334        );
2335
2336        // The bytes the store was handed must account for the whole chunk,
2337        // dictionary page included. Holding the dictionary page apart from the
2338        // store would make this fall short by the dictionary page's size.
2339        assert_eq!(
2340            bytes_put.load(std::sync::atomic::Ordering::Relaxed) as i64,
2341            column.compressed_size(),
2342            "the dictionary page must pass through the store like any other page"
2343        );
2344    }
2345
2346    #[test]
2347    fn arrow_writer() {
2348        // define schema
2349        let schema = Schema::new(vec![
2350            Field::new("a", DataType::Int32, false),
2351            Field::new("b", DataType::Int32, true),
2352        ]);
2353
2354        // create some data
2355        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2356        let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2357
2358        // build a record batch
2359        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a), Arc::new(b)]).unwrap();
2360
2361        roundtrip(batch, Some(SMALL_SIZE / 2));
2362    }
2363
2364    fn get_bytes_after_close(schema: SchemaRef, expected_batch: &RecordBatch) -> Vec<u8> {
2365        let mut buffer = vec![];
2366
2367        let mut writer = ArrowWriter::try_new(&mut buffer, schema, None).unwrap();
2368        writer.write(expected_batch).unwrap();
2369        writer.close().unwrap();
2370
2371        buffer
2372    }
2373
2374    fn get_bytes_by_into_inner(schema: SchemaRef, expected_batch: &RecordBatch) -> Vec<u8> {
2375        let mut writer = ArrowWriter::try_new(Vec::new(), schema, None).unwrap();
2376        writer.write(expected_batch).unwrap();
2377        writer.into_inner().unwrap()
2378    }
2379
2380    #[test]
2381    fn roundtrip_bytes() {
2382        // define schema
2383        let schema = Arc::new(Schema::new(vec![
2384            Field::new("a", DataType::Int32, false),
2385            Field::new("b", DataType::Int32, true),
2386        ]));
2387
2388        // create some data
2389        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2390        let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2391
2392        // build a record batch
2393        let expected_batch =
2394            RecordBatch::try_new(schema.clone(), vec![Arc::new(a), Arc::new(b)]).unwrap();
2395
2396        for buffer in [
2397            get_bytes_after_close(schema.clone(), &expected_batch),
2398            get_bytes_by_into_inner(schema, &expected_batch),
2399        ] {
2400            let cursor = Bytes::from(buffer);
2401            let mut record_batch_reader = ParquetRecordBatchReader::try_new(cursor, 1024).unwrap();
2402
2403            let actual_batch = record_batch_reader
2404                .next()
2405                .expect("No batch found")
2406                .expect("Unable to get batch");
2407
2408            assert_eq!(expected_batch.schema(), actual_batch.schema());
2409            assert_eq!(expected_batch.num_columns(), actual_batch.num_columns());
2410            assert_eq!(expected_batch.num_rows(), actual_batch.num_rows());
2411            for i in 0..expected_batch.num_columns() {
2412                let expected_data = expected_batch.column(i).to_data();
2413                let actual_data = actual_batch.column(i).to_data();
2414
2415                assert_eq!(expected_data, actual_data);
2416            }
2417        }
2418    }
2419
2420    #[test]
2421    #[cfg_attr(miri, ignore)] // Takes too long
2422    fn arrow_writer_non_null() {
2423        let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
2424        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2425
2426        RoundTripTest::new(Arc::new(a))
2427            .with_schema(Arc::new(schema))
2428            .run();
2429    }
2430
2431    #[test]
2432    #[cfg_attr(miri, ignore)] // Takes too long
2433    fn arrow_writer_list() {
2434        // define schema
2435        let schema = Schema::new(vec![Field::new(
2436            "a",
2437            DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false))),
2438            true,
2439        )]);
2440
2441        // create some data
2442        let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
2443
2444        // Construct a buffer for value offsets, for the nested array:
2445        //  [[1], [2, 3], null, [4, 5, 6], [7, 8, 9, 10]]
2446        let a_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice());
2447
2448        // Construct a list array from the above two
2449        let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
2450            DataType::Int32,
2451            false,
2452        ))))
2453        .len(5)
2454        .add_buffer(a_value_offsets)
2455        .add_child_data(a_values.into_data())
2456        .null_bit_buffer(Some(Buffer::from([0b00011011])))
2457        .build()
2458        .unwrap();
2459        let a = ListArray::from(a_list_data);
2460        assert_eq!(a.null_count(), 1);
2461
2462        RoundTripTest::new(Arc::new(a))
2463            .with_schema(Arc::new(schema))
2464            .run();
2465    }
2466
2467    #[test]
2468    #[cfg_attr(miri, ignore)] // Takes too long
2469    fn arrow_writer_list_non_null() {
2470        // define schema
2471        let schema = Schema::new(vec![Field::new(
2472            "a",
2473            DataType::List(Arc::new(Field::new_list_field(DataType::Int32, false))),
2474            false,
2475        )]);
2476
2477        // create some data
2478        let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
2479
2480        // Construct a buffer for value offsets, for the nested array:
2481        //  [[1], [2, 3], [], [4, 5, 6], [7, 8, 9, 10]]
2482        let a_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice());
2483
2484        // Construct a list array from the above two
2485        let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
2486            DataType::Int32,
2487            false,
2488        ))))
2489        .len(5)
2490        .add_buffer(a_value_offsets)
2491        .add_child_data(a_values.into_data())
2492        .build()
2493        .unwrap();
2494        let a = ListArray::from(a_list_data);
2495        assert_eq!(a.null_count(), 0);
2496
2497        RoundTripTest::new(Arc::new(a))
2498            .with_schema(Arc::new(schema))
2499            .run();
2500    }
2501
2502    #[test]
2503    #[cfg_attr(miri, ignore)] // Takes too long
2504    fn arrow_writer_list_view() {
2505        let list_field = Arc::new(Field::new_list_field(DataType::Int32, false));
2506        let schema = Schema::new(vec![Field::new(
2507            "a",
2508            DataType::ListView(list_field.clone()),
2509            true,
2510        )]);
2511
2512        //  [[1], [2, 3], null, [4, 5, 6], [7, 8, 9, 10]]
2513        let a = ListViewArray::new(
2514            list_field,
2515            vec![0, 1, 0, 3, 6].into(),
2516            vec![1, 2, 0, 3, 4].into(),
2517            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
2518            Some(vec![true, true, false, true, true].into()),
2519        );
2520        assert_eq!(a.null_count(), 1);
2521
2522        RoundTripTest::new(Arc::new(a))
2523            .with_schema(Arc::new(schema))
2524            .run();
2525    }
2526
2527    #[test]
2528    #[cfg_attr(miri, ignore)] // Takes too long
2529    fn arrow_writer_list_view_non_null() {
2530        let list_field = Arc::new(Field::new_list_field(DataType::Int32, false));
2531        let schema = Schema::new(vec![Field::new(
2532            "a",
2533            DataType::ListView(list_field.clone()),
2534            false,
2535        )]);
2536
2537        //  [[1], [2, 3], [], [4, 5, 6], [7, 8, 9, 10]]
2538        let a = ListViewArray::new(
2539            list_field,
2540            vec![0, 1, 0, 3, 6].into(),
2541            vec![1, 2, 0, 3, 4].into(),
2542            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
2543            None,
2544        );
2545        assert_eq!(a.null_count(), 0);
2546
2547        RoundTripTest::new(Arc::new(a))
2548            .with_schema(Arc::new(schema))
2549            .run();
2550    }
2551
2552    #[test]
2553    #[cfg_attr(miri, ignore)] // Takes too long
2554    fn arrow_writer_list_view_out_of_order() {
2555        let list_field = Arc::new(Field::new_list_field(DataType::Int32, false));
2556        let schema = Schema::new(vec![Field::new(
2557            "a",
2558            DataType::ListView(list_field.clone()),
2559            false,
2560        )]);
2561
2562        // [[1], [2, 3], [], [7, 8, 9, 10], [4, 5, 6]] - out of order offsets
2563        let a = ListViewArray::new(
2564            list_field,
2565            vec![0, 1, 0, 6, 3].into(),
2566            vec![1, 2, 0, 4, 3].into(),
2567            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
2568            None,
2569        );
2570        assert_eq!(a.null_count(), 0);
2571
2572        RoundTripTest::new(Arc::new(a))
2573            .with_schema(Arc::new(schema))
2574            .run();
2575    }
2576
2577    #[test]
2578    #[cfg_attr(miri, ignore)] // Takes too long
2579    fn arrow_writer_large_list_view() {
2580        let list_field = Arc::new(Field::new_list_field(DataType::Int32, false));
2581        let schema = Schema::new(vec![Field::new(
2582            "a",
2583            DataType::LargeListView(list_field.clone()),
2584            true,
2585        )]);
2586
2587        //  [[1], [2, 3], null, [4, 5, 6], [7, 8, 9, 10]]
2588        let a = LargeListViewArray::new(
2589            list_field,
2590            vec![0i64, 1, 0, 3, 6].into(),
2591            vec![1i64, 2, 0, 3, 4].into(),
2592            Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])),
2593            Some(vec![true, true, false, true, true].into()),
2594        );
2595        assert_eq!(a.null_count(), 1);
2596
2597        RoundTripTest::new(Arc::new(a))
2598            .with_schema(Arc::new(schema))
2599            .run();
2600    }
2601
2602    #[test]
2603    #[cfg_attr(miri, ignore)] // Takes too long
2604    fn arrow_writer_list_view_with_struct() {
2605        // Test ListView containing Struct: ListView<Struct<Int32, Utf8>>
2606        let struct_fields = Fields::from(vec![
2607            Field::new("id", DataType::Int32, false),
2608            Field::new("name", DataType::Utf8, false),
2609        ]);
2610        let struct_type = DataType::Struct(struct_fields.clone());
2611        let list_field = Arc::new(Field::new("item", struct_type.clone(), false));
2612
2613        let schema = Schema::new(vec![Field::new(
2614            "a",
2615            DataType::ListView(list_field.clone()),
2616            true,
2617        )]);
2618
2619        // Create struct values
2620        let id_array = Int32Array::from(vec![1, 2, 3, 4, 5]);
2621        let name_array = StringArray::from(vec!["a", "b", "c", "d", "e"]);
2622        let struct_array = StructArray::new(
2623            struct_fields,
2624            vec![Arc::new(id_array), Arc::new(name_array)],
2625            None,
2626        );
2627
2628        // Create ListView: [{1, "a"}, {2, "b"}], null, [{3, "c"}, {4, "d"}, {5, "e"}]
2629        let list_view = ListViewArray::new(
2630            list_field,
2631            vec![0, 2, 2].into(), // offsets
2632            vec![2, 0, 3].into(), // sizes
2633            Arc::new(struct_array),
2634            Some(vec![true, false, true].into()),
2635        );
2636        assert_eq!(list_view.null_count(), 1);
2637
2638        RoundTripTest::new(Arc::new(list_view))
2639            .with_schema(Arc::new(schema))
2640            .run();
2641    }
2642
2643    #[test]
2644    #[cfg_attr(miri, ignore)] // Takes too long
2645    fn arrow_writer_binary() {
2646        let raw_string_values = vec!["foo", "bar", "baz", "quux"];
2647        let raw_binary_values = [
2648            b"foo".to_vec(),
2649            b"bar".to_vec(),
2650            b"baz".to_vec(),
2651            b"quux".to_vec(),
2652        ];
2653        let raw_binary_value_refs = raw_binary_values
2654            .iter()
2655            .map(|x| x.as_slice())
2656            .collect::<Vec<_>>();
2657
2658        let string_values = StringArray::from(raw_string_values.clone());
2659        let binary_values = BinaryArray::from(raw_binary_value_refs);
2660        assert_eq!(string_values.null_count(), 0);
2661        assert_eq!(binary_values.null_count(), 0);
2662
2663        RoundTripTest::new(Arc::new(string_values)).run();
2664        RoundTripTest::new(Arc::new(binary_values)).run();
2665    }
2666
2667    #[test]
2668    #[cfg_attr(miri, ignore)] // Takes too long
2669    fn arrow_writer_binary_view() {
2670        let raw_string_values = vec!["foo", "bar", "large payload over 12 bytes", "lulu"];
2671        let raw_binary_values = vec![
2672            b"foo".to_vec(),
2673            b"bar".to_vec(),
2674            b"large payload over 12 bytes".to_vec(),
2675            b"lulu".to_vec(),
2676        ];
2677        let nullable_string_values =
2678            vec![Some("foo"), None, Some("large payload over 12 bytes"), None];
2679
2680        let string_view_values = StringViewArray::from(raw_string_values);
2681        let binary_view_values = BinaryViewArray::from_iter_values(raw_binary_values);
2682        let nullable_string_view_values = StringViewArray::from(nullable_string_values);
2683
2684        RoundTripTest::new(Arc::new(string_view_values)).run();
2685        RoundTripTest::new(Arc::new(binary_view_values)).run();
2686        RoundTripTest::new(Arc::new(nullable_string_view_values)).run();
2687    }
2688
2689    #[test]
2690    #[cfg_attr(miri, ignore)] // Takes too long
2691    fn arrow_writer_binary_view_long_value() {
2692        // There is special case validation for long values (greater than 128)
2693        // 128 encodes as 0x80 0x00 0x00 0x00 in little endian, which should
2694        // trigger the long-string UTF-8 validation branch in the plain decoder.
2695        let long = "a".repeat(128);
2696        let raw_string_values = vec!["foo", long.as_str(), "bar"];
2697        let raw_binary_values = vec![b"foo".to_vec(), long.as_bytes().to_vec(), b"bar".to_vec()];
2698
2699        let string_view_values: ArrayRef = Arc::new(StringViewArray::from(raw_string_values));
2700        let binary_view_values: ArrayRef =
2701            Arc::new(BinaryViewArray::from_iter_values(raw_binary_values));
2702
2703        RoundTripTest::new(Arc::clone(&string_view_values))
2704            .with_nullable(false)
2705            .run();
2706        RoundTripTest::new(Arc::clone(&binary_view_values))
2707            .with_nullable(false)
2708            .run();
2709    }
2710
2711    /// Test round-trip of Dictionary<UInt32, Utf8View> and
2712    /// Dictionary<UInt32, BinaryView> typed columns.
2713    #[test]
2714    fn arrow_writer_string_view_dictionary() {
2715        let raw_string_values = vec!["a", "b", "large payload over 12 bytes"];
2716        let raw_binary_values = vec![
2717            b"a".to_vec(),
2718            b"b".to_vec(),
2719            b"large payload over 12 bytes".to_vec(),
2720        ];
2721
2722        let keys = UInt32Array::from(vec![Some(0), None, Some(2), Some(1), None]);
2723
2724        let string_view_values = Arc::new(StringViewArray::from(raw_string_values));
2725        let string_dict: ArrayRef = Arc::new(
2726            DictionaryArray::<UInt32Type>::try_new(keys.clone(), string_view_values).unwrap(),
2727        );
2728
2729        let binary_view_values = Arc::new(BinaryViewArray::from_iter_values(raw_binary_values));
2730        let binary_dict: ArrayRef =
2731            Arc::new(DictionaryArray::<UInt32Type>::try_new(keys, binary_view_values).unwrap());
2732
2733        RoundTripTest::new(string_dict).run();
2734        RoundTripTest::new(binary_dict).run();
2735    }
2736
2737    fn get_decimal_batch(precision: u8, scale: i8) -> RecordBatch {
2738        let decimal_field = Field::new("a", DataType::Decimal128(precision, scale), false);
2739        let schema = Schema::new(vec![decimal_field]);
2740
2741        let decimal_values = vec![10_000, 50_000, 0, -100]
2742            .into_iter()
2743            .map(Some)
2744            .collect::<Decimal128Array>()
2745            .with_precision_and_scale(precision, scale)
2746            .unwrap();
2747
2748        RecordBatch::try_new(Arc::new(schema), vec![Arc::new(decimal_values)]).unwrap()
2749    }
2750
2751    #[test]
2752    fn arrow_writer_decimal() {
2753        // int32 to store the decimal value
2754        let batch_int32_decimal = get_decimal_batch(5, 2);
2755        roundtrip(batch_int32_decimal, Some(SMALL_SIZE / 2));
2756        // int64 to store the decimal value
2757        let batch_int64_decimal = get_decimal_batch(12, 2);
2758        roundtrip(batch_int64_decimal, Some(SMALL_SIZE / 2));
2759        // fixed_length_byte_array to store the decimal value
2760        let batch_fixed_len_byte_array_decimal = get_decimal_batch(30, 2);
2761        roundtrip(batch_fixed_len_byte_array_decimal, Some(SMALL_SIZE / 2));
2762    }
2763
2764    #[test]
2765    #[cfg_attr(miri, ignore)] // Takes too long
2766    fn arrow_writer_complex() {
2767        // define schema
2768        let struct_field_d = Arc::new(Field::new("d", DataType::Float64, true));
2769        let struct_field_f = Arc::new(Field::new("f", DataType::Float32, true));
2770        let struct_field_g = Arc::new(Field::new_list(
2771            "g",
2772            Field::new_list_field(DataType::Int16, true),
2773            false,
2774        ));
2775        let struct_field_h = Arc::new(Field::new_list(
2776            "h",
2777            Field::new_list_field(DataType::Int16, false),
2778            true,
2779        ));
2780        let struct_field_e = Arc::new(Field::new_struct(
2781            "e",
2782            vec![
2783                struct_field_f.clone(),
2784                struct_field_g.clone(),
2785                struct_field_h.clone(),
2786            ],
2787            false,
2788        ));
2789        let schema = Schema::new(vec![
2790            Field::new("a", DataType::Int32, false),
2791            Field::new("b", DataType::Int32, true),
2792            Field::new_struct(
2793                "c",
2794                vec![struct_field_d.clone(), struct_field_e.clone()],
2795                false,
2796            ),
2797        ]);
2798
2799        // create some data
2800        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
2801        let b = Int32Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2802        let d = Float64Array::from(vec![None, None, None, Some(1.0), None]);
2803        let f = Float32Array::from(vec![Some(0.0), None, Some(333.3), None, Some(5.25)]);
2804
2805        let g_value = Int16Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
2806
2807        // Construct a buffer for value offsets, for the nested array:
2808        //  [[1], [2, 3], [], [4, 5, 6], [7, 8, 9, 10]]
2809        let g_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice());
2810
2811        // Construct a list array from the above two
2812        let g_list_data = ArrayData::builder(struct_field_g.data_type().clone())
2813            .len(5)
2814            .add_buffer(g_value_offsets.clone())
2815            .add_child_data(g_value.to_data())
2816            .build()
2817            .unwrap();
2818        let g = ListArray::from(g_list_data);
2819        // The difference between g and h is that h has a null bitmap
2820        let h_list_data = ArrayData::builder(struct_field_h.data_type().clone())
2821            .len(5)
2822            .add_buffer(g_value_offsets)
2823            .add_child_data(g_value.to_data())
2824            .null_bit_buffer(Some(Buffer::from([0b00011011])))
2825            .build()
2826            .unwrap();
2827        let h = ListArray::from(h_list_data);
2828
2829        let e = StructArray::from(vec![
2830            (struct_field_f, Arc::new(f) as ArrayRef),
2831            (struct_field_g, Arc::new(g) as ArrayRef),
2832            (struct_field_h, Arc::new(h) as ArrayRef),
2833        ]);
2834
2835        let c = StructArray::from(vec![
2836            (struct_field_d, Arc::new(d) as ArrayRef),
2837            (struct_field_e, Arc::new(e) as ArrayRef),
2838        ]);
2839
2840        // build a record batch
2841        let batch = RecordBatch::try_new(
2842            Arc::new(schema),
2843            vec![Arc::new(a), Arc::new(b), Arc::new(c)],
2844        )
2845        .unwrap();
2846
2847        roundtrip(batch.clone(), Some(SMALL_SIZE / 2));
2848        roundtrip(batch, Some(SMALL_SIZE / 3));
2849    }
2850
2851    #[test]
2852    fn arrow_writer_complex_mixed() {
2853        // This test was added while investigating https://github.com/apache/arrow-rs/issues/244.
2854        // It was subsequently fixed while investigating https://github.com/apache/arrow-rs/issues/245.
2855
2856        // define schema
2857        let offset_field = Arc::new(Field::new("offset", DataType::Int32, false));
2858        let partition_field = Arc::new(Field::new("partition", DataType::Int64, true));
2859        let topic_field = Arc::new(Field::new("topic", DataType::Utf8, true));
2860        let schema = Schema::new(vec![Field::new(
2861            "some_nested_object",
2862            DataType::Struct(Fields::from(vec![
2863                offset_field.clone(),
2864                partition_field.clone(),
2865                topic_field.clone(),
2866            ])),
2867            false,
2868        )]);
2869
2870        // create some data
2871        let offset = Int32Array::from(vec![1, 2, 3, 4, 5]);
2872        let partition = Int64Array::from(vec![Some(1), None, None, Some(4), Some(5)]);
2873        let topic = StringArray::from(vec![Some("A"), None, Some("A"), Some(""), None]);
2874
2875        let some_nested_object = StructArray::from(vec![
2876            (offset_field, Arc::new(offset) as ArrayRef),
2877            (partition_field, Arc::new(partition) as ArrayRef),
2878            (topic_field, Arc::new(topic) as ArrayRef),
2879        ]);
2880
2881        // build a record batch
2882        let batch =
2883            RecordBatch::try_new(Arc::new(schema), vec![Arc::new(some_nested_object)]).unwrap();
2884
2885        roundtrip(batch, Some(SMALL_SIZE / 2));
2886    }
2887
2888    #[test]
2889    fn arrow_writer_map() {
2890        // Note: we are using the JSON Arrow reader for brevity
2891        let json_content = r#"
2892        {"stocks":{"long": "$AAA", "short": "$BBB"}}
2893        {"stocks":{"long": null, "long": "$CCC", "short": null}}
2894        {"stocks":{"hedged": "$YYY", "long": null, "short": "$D"}}
2895        "#;
2896        let entries_struct_type = DataType::Struct(Fields::from(vec![
2897            Field::new(Field::MAP_KEY_FIELD_DEFAULT_NAME, DataType::Utf8, false),
2898            Field::new(Field::MAP_VALUE_FIELD_DEFAULT_NAME, DataType::Utf8, true),
2899        ]));
2900        let stocks_field = Field::new(
2901            "stocks",
2902            DataType::Map(
2903                Arc::new(Field::new(
2904                    Field::MAP_ENTRIES_FIELD_DEFAULT_NAME,
2905                    entries_struct_type,
2906                    false,
2907                )),
2908                false,
2909            ),
2910            true,
2911        );
2912        let schema = Arc::new(Schema::new(vec![stocks_field]));
2913        let builder = arrow::json::ReaderBuilder::new(schema).with_batch_size(64);
2914        let mut reader = builder.build(std::io::Cursor::new(json_content)).unwrap();
2915
2916        let batch = reader.next().unwrap().unwrap();
2917        roundtrip(batch, None);
2918    }
2919
2920    #[test]
2921    fn arrow_writer_2_level_struct() {
2922        // tests writing <struct<struct<primitive>>
2923        let field_c = Field::new("c", DataType::Int32, true);
2924        let field_b = Field::new("b", DataType::Struct(vec![field_c].into()), true);
2925        let type_a = DataType::Struct(vec![field_b.clone()].into());
2926        let field_a = Field::new("a", type_a, true);
2927        let schema = Schema::new(vec![field_a.clone()]);
2928
2929        // create data
2930        let c = Int32Array::from(vec![Some(1), None, Some(3), None, None, Some(6)]);
2931        let b_data = ArrayDataBuilder::new(field_b.data_type().clone())
2932            .len(6)
2933            .null_bit_buffer(Some(Buffer::from([0b00100111])))
2934            .add_child_data(c.into_data())
2935            .build()
2936            .unwrap();
2937        let b = StructArray::from(b_data);
2938        let a_data = ArrayDataBuilder::new(field_a.data_type().clone())
2939            .len(6)
2940            .null_bit_buffer(Some(Buffer::from([0b00101111])))
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(), 1);
2947        assert_eq!(a.column(0).null_count(), 2);
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_non_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(), false);
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            .add_child_data(c.into_data())
2970            .build()
2971            .unwrap();
2972        let b = StructArray::from(b_data);
2973        let a_data = ArrayDataBuilder::new(type_a)
2974            .len(6)
2975            .add_child_data(b.into_data())
2976            .build()
2977            .unwrap();
2978        let a = StructArray::from(a_data);
2979
2980        assert_eq!(a.null_count(), 0);
2981        assert_eq!(a.column(0).null_count(), 0);
2982
2983        // build a racord batch
2984        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
2985
2986        roundtrip(batch, Some(SMALL_SIZE / 2));
2987    }
2988
2989    #[test]
2990    fn arrow_writer_2_level_struct_mixed_null() {
2991        // tests writing <struct<struct<primitive>>
2992        let field_c = Field::new("c", DataType::Int32, false);
2993        let type_b = DataType::Struct(vec![field_c].into());
2994        let field_b = Field::new("b", type_b.clone(), true);
2995        let type_a = DataType::Struct(vec![field_b].into());
2996        let field_a = Field::new("a", type_a.clone(), false);
2997        let schema = Schema::new(vec![field_a]);
2998
2999        // create data
3000        let c = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
3001        let b_data = ArrayDataBuilder::new(type_b)
3002            .len(6)
3003            .null_bit_buffer(Some(Buffer::from([0b00100111])))
3004            .add_child_data(c.into_data())
3005            .build()
3006            .unwrap();
3007        let b = StructArray::from(b_data);
3008        // a intentionally has no null buffer, to test that this is handled correctly
3009        let a_data = ArrayDataBuilder::new(type_a)
3010            .len(6)
3011            .add_child_data(b.into_data())
3012            .build()
3013            .unwrap();
3014        let a = StructArray::from(a_data);
3015
3016        assert_eq!(a.null_count(), 0);
3017        assert_eq!(a.column(0).null_count(), 2);
3018
3019        // build a racord batch
3020        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
3021
3022        roundtrip(batch, Some(SMALL_SIZE / 2));
3023    }
3024
3025    #[test]
3026    fn arrow_writer_2_level_struct_mixed_null_2() {
3027        // tests writing <struct<struct<primitive>>, where the primitive columns are non-null.
3028        let field_c = Field::new("c", DataType::Int32, false);
3029        let field_d = Field::new("d", DataType::FixedSizeBinary(4), false);
3030        let field_e = Field::new(
3031            "e",
3032            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
3033            false,
3034        );
3035
3036        let field_b = Field::new(
3037            "b",
3038            DataType::Struct(vec![field_c, field_d, field_e].into()),
3039            false,
3040        );
3041        let type_a = DataType::Struct(vec![field_b.clone()].into());
3042        let field_a = Field::new("a", type_a, true);
3043        let schema = Schema::new(vec![field_a.clone()]);
3044
3045        // create data
3046        let c = Int32Array::from_iter_values(0..6);
3047        let d = FixedSizeBinaryArray::try_from_iter(
3048            ["aaaa", "bbbb", "cccc", "dddd", "eeee", "ffff"].into_iter(),
3049        )
3050        .expect("four byte values");
3051        let e = Int32DictionaryArray::from_iter(["one", "two", "three", "four", "five", "one"]);
3052        let b_data = ArrayDataBuilder::new(field_b.data_type().clone())
3053            .len(6)
3054            .add_child_data(c.into_data())
3055            .add_child_data(d.into_data())
3056            .add_child_data(e.into_data())
3057            .build()
3058            .unwrap();
3059        let b = StructArray::from(b_data);
3060        let a_data = ArrayDataBuilder::new(field_a.data_type().clone())
3061            .len(6)
3062            .null_bit_buffer(Some(Buffer::from([0b00100101])))
3063            .add_child_data(b.into_data())
3064            .build()
3065            .unwrap();
3066        let a = StructArray::from(a_data);
3067
3068        assert_eq!(a.null_count(), 3);
3069        assert_eq!(a.column(0).null_count(), 0);
3070
3071        // build a record batch
3072        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
3073
3074        roundtrip(batch, Some(SMALL_SIZE / 2));
3075    }
3076
3077    #[test]
3078    fn test_fixed_size_binary_in_dict() {
3079        fn test_fixed_size_binary_in_dict_inner<K>()
3080        where
3081            K: ArrowDictionaryKeyType,
3082            K::Native: FromPrimitive + ToPrimitive + TryFrom<u8>,
3083            <<K as arrow_array::ArrowPrimitiveType>::Native as TryFrom<u8>>::Error: std::fmt::Debug,
3084        {
3085            let field = Field::new(
3086                "a",
3087                DataType::Dictionary(
3088                    Box::new(K::DATA_TYPE),
3089                    Box::new(DataType::FixedSizeBinary(4)),
3090                ),
3091                false,
3092            );
3093            let schema = Schema::new(vec![field]);
3094
3095            let keys: Vec<K::Native> = vec![
3096                K::Native::try_from(0u8).unwrap(),
3097                K::Native::try_from(0u8).unwrap(),
3098                K::Native::try_from(1u8).unwrap(),
3099            ];
3100            let keys = PrimitiveArray::<K>::from_iter_values(keys);
3101            let values = FixedSizeBinaryArray::try_from_iter(
3102                vec![vec![0, 0, 0, 0], vec![1, 1, 1, 1]].into_iter(),
3103            )
3104            .unwrap();
3105
3106            let data = DictionaryArray::<K>::new(keys, Arc::new(values));
3107            let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(data)]).unwrap();
3108            roundtrip(batch, None);
3109        }
3110
3111        test_fixed_size_binary_in_dict_inner::<UInt8Type>();
3112        test_fixed_size_binary_in_dict_inner::<UInt16Type>();
3113        test_fixed_size_binary_in_dict_inner::<UInt32Type>();
3114        test_fixed_size_binary_in_dict_inner::<UInt16Type>();
3115        test_fixed_size_binary_in_dict_inner::<Int8Type>();
3116        test_fixed_size_binary_in_dict_inner::<Int16Type>();
3117        test_fixed_size_binary_in_dict_inner::<Int32Type>();
3118        test_fixed_size_binary_in_dict_inner::<Int64Type>();
3119    }
3120
3121    #[test]
3122    fn test_empty_dict() {
3123        let struct_fields = Fields::from(vec![Field::new(
3124            "dict",
3125            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
3126            false,
3127        )]);
3128
3129        let schema = Schema::new(vec![Field::new_struct(
3130            "struct",
3131            struct_fields.clone(),
3132            true,
3133        )]);
3134        let dictionary = Arc::new(DictionaryArray::new(
3135            Int32Array::new_null(5),
3136            Arc::new(StringArray::new_null(0)),
3137        ));
3138
3139        let s = StructArray::new(
3140            struct_fields,
3141            vec![dictionary],
3142            Some(NullBuffer::new_null(5)),
3143        );
3144
3145        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(s)]).unwrap();
3146        roundtrip(batch, None);
3147    }
3148    #[test]
3149    fn arrow_writer_page_size() {
3150        let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)]));
3151
3152        let mut builder = StringBuilder::with_capacity(100, 329 * 10_000);
3153
3154        // Generate an array of 10 unique 10 character string
3155        for i in 0..10 {
3156            let value = i
3157                .to_string()
3158                .repeat(10)
3159                .chars()
3160                .take(10)
3161                .collect::<String>();
3162
3163            builder.append_value(value);
3164        }
3165
3166        let array = Arc::new(builder.finish());
3167
3168        let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
3169
3170        let file = tempfile::tempfile().unwrap();
3171
3172        // Set everything very low so we fallback to PLAIN encoding after the first row
3173        let props = WriterProperties::builder()
3174            .set_data_page_size_limit(1)
3175            .set_dictionary_page_size_limit(1)
3176            .set_write_batch_size(1)
3177            .build();
3178
3179        let mut writer =
3180            ArrowWriter::try_new(file.try_clone().unwrap(), batch.schema(), Some(props))
3181                .expect("Unable to write file");
3182        writer.write(&batch).unwrap();
3183        writer.close().unwrap();
3184
3185        let options = ReadOptionsBuilder::new().with_page_index().build();
3186        let reader =
3187            SerializedFileReader::new_with_options(file.try_clone().unwrap(), options).unwrap();
3188
3189        let column = reader.metadata().row_group(0).columns();
3190
3191        assert_eq!(column.len(), 1);
3192
3193        // We should write one row before falling back to PLAIN encoding so there should still be a
3194        // dictionary page.
3195        assert!(
3196            column[0].dictionary_page_offset().is_some(),
3197            "Expected a dictionary page"
3198        );
3199
3200        let page_index = reader
3201            .metadata()
3202            .page_index()
3203            .expect("page index should be present");
3204        let page_locations = page_index
3205            .page_locations(0, 0)
3206            .expect("page locations should exist");
3207
3208        // We should fallback to PLAIN encoding after the first row and our max page size is 1 bytes
3209        // so we expect one dictionary encoded page and then a page per row thereafter.
3210        assert_eq!(
3211            page_locations.len(),
3212            10,
3213            "Expected 10 pages but got {page_locations:#?}"
3214        );
3215    }
3216
3217    #[test]
3218    #[cfg_attr(miri, ignore)] // inline assembly is not supported
3219    fn arrow_writer_float_nans() {
3220        let f16_field = Field::new("a", DataType::Float16, false);
3221        let f32_field = Field::new("b", DataType::Float32, false);
3222        let f64_field = Field::new("c", DataType::Float64, false);
3223        let schema = Schema::new(vec![f16_field, f32_field, f64_field]);
3224
3225        let f16_values = (0..MEDIUM_SIZE)
3226            .map(|i| {
3227                Some(if i % 2 == 0 {
3228                    f16::NAN
3229                } else {
3230                    f16::from_f32(i as f32)
3231                })
3232            })
3233            .collect::<Float16Array>();
3234
3235        let f32_values = (0..MEDIUM_SIZE)
3236            .map(|i| Some(if i % 2 == 0 { f32::NAN } else { i as f32 }))
3237            .collect::<Float32Array>();
3238
3239        let f64_values = (0..MEDIUM_SIZE)
3240            .map(|i| Some(if i % 2 == 0 { f64::NAN } else { i as f64 }))
3241            .collect::<Float64Array>();
3242
3243        let batch = RecordBatch::try_new(
3244            Arc::new(schema),
3245            vec![
3246                Arc::new(f16_values),
3247                Arc::new(f32_values),
3248                Arc::new(f64_values),
3249            ],
3250        )
3251        .unwrap();
3252
3253        roundtrip(batch, None);
3254    }
3255
3256    const SMALL_SIZE: usize = 7;
3257    const MEDIUM_SIZE: usize = 63;
3258
3259    // Write the batch to parquet and read it back out, ensuring
3260    // that what comes out is the same as what was written in
3261    fn roundtrip(expected_batch: RecordBatch, max_row_group_size: Option<usize>) -> Vec<Bytes> {
3262        let mut files = vec![];
3263        for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] {
3264            let mut props = WriterProperties::builder().set_writer_version(version);
3265
3266            if let Some(size) = max_row_group_size {
3267                props = props.set_max_row_group_row_count(Some(size))
3268            }
3269
3270            let props = props.build();
3271            files.push(roundtrip_opts(&expected_batch, props))
3272        }
3273        files
3274    }
3275
3276    // Round trip the specified record batch with the specified writer properties,
3277    // to an in-memory file, and validate the arrays using the specified function.
3278    // Returns the in-memory file.
3279    fn roundtrip_opts_with_array_validation<F>(
3280        expected_batch: &RecordBatch,
3281        props: WriterProperties,
3282        validate: F,
3283    ) -> Bytes
3284    where
3285        F: Fn(&ArrayData, &ArrayData),
3286    {
3287        let mut file = vec![];
3288
3289        let mut writer = ArrowWriter::try_new(&mut file, expected_batch.schema(), Some(props))
3290            .expect("Unable to write file");
3291        writer.write(expected_batch).unwrap();
3292        writer.close().unwrap();
3293
3294        let file = Bytes::from(file);
3295        let mut record_batch_reader =
3296            ParquetRecordBatchReader::try_new(file.clone(), 1024).unwrap();
3297
3298        let actual_batch = record_batch_reader
3299            .next()
3300            .expect("No batch found")
3301            .expect("Unable to get batch");
3302
3303        assert_eq!(expected_batch.schema(), actual_batch.schema());
3304        assert_eq!(expected_batch.num_columns(), actual_batch.num_columns());
3305        assert_eq!(expected_batch.num_rows(), actual_batch.num_rows());
3306        for i in 0..expected_batch.num_columns() {
3307            let expected_data = expected_batch.column(i).to_data();
3308            let actual_data = actual_batch.column(i).to_data();
3309            validate(&expected_data, &actual_data);
3310        }
3311
3312        file
3313    }
3314
3315    fn roundtrip_opts(expected_batch: &RecordBatch, props: WriterProperties) -> Bytes {
3316        roundtrip_opts_with_array_validation(expected_batch, props, |a, b| {
3317            a.validate_full().expect("valid expected data");
3318            b.validate_full().expect("valid actual data");
3319            assert_eq!(a, b)
3320        })
3321    }
3322
3323    /// Round trip testing fixture:
3324    ///
3325    /// Tests based on this fixture write data to parquet and then read it back.
3326    struct RoundTripTest {
3327        values: ArrayRef,
3328        /// Optionally supplied schema
3329        schema: Option<SchemaRef>,
3330        /// If the created schema should be nullable. Defaults to true. Ignored
3331        /// if schema is set to Some.
3332        nullable: bool,
3333        bloom_filter: bool,
3334        bloom_filter_ndv: Option<u64>,
3335        bloom_filter_position: BloomFilterPosition,
3336    }
3337
3338    impl RoundTripTest {
3339        /// Create a test for round tripping values with a nullable schema
3340        fn new(values: ArrayRef) -> Self {
3341            Self {
3342                values,
3343                schema: None,
3344                nullable: true,
3345                bloom_filter: false,
3346                bloom_filter_ndv: None,
3347                bloom_filter_position: BloomFilterPosition::AfterRowGroup,
3348            }
3349        }
3350
3351        /// Set the schema
3352        fn with_schema(mut self, schema: SchemaRef) -> Self {
3353            self.schema = Some(schema);
3354            self
3355        }
3356
3357        /// Set the nullable flag
3358        fn with_nullable(mut self, nullable: bool) -> Self {
3359            self.nullable = nullable;
3360            self
3361        }
3362
3363        /// Set bloom filter
3364        fn with_bloom_filter(mut self, bloom_filter: bool) -> Self {
3365            self.bloom_filter = bloom_filter;
3366            self
3367        }
3368
3369        /// Set bloom filter max ndv
3370        fn with_bloom_filter_ndv(mut self, bloom_filter_ndv: u64) -> Self {
3371            self.bloom_filter_ndv = Some(bloom_filter_ndv);
3372            self
3373        }
3374
3375        /// Set bloom filter position
3376        fn with_bloom_filter_position(
3377            mut self,
3378            bloom_filter_position: BloomFilterPosition,
3379        ) -> Self {
3380            self.bloom_filter_position = bloom_filter_position;
3381            self
3382        }
3383
3384        /// Run the test specified by the options, returning the encoded Parquet bytes
3385        fn run(self) -> Vec<Bytes> {
3386            let RoundTripTest {
3387                values,
3388                schema,
3389                nullable,
3390                bloom_filter,
3391                bloom_filter_ndv,
3392                bloom_filter_position,
3393            } = self;
3394
3395            let schema = schema.unwrap_or_else(|| {
3396                let data_type = values.data_type().clone();
3397                Arc::new(Schema::new(vec![Field::new("col", data_type, nullable)]))
3398            });
3399
3400            let encodings = match values.data_type() {
3401                DataType::Utf8 | DataType::LargeUtf8 | DataType::Binary | DataType::LargeBinary => {
3402                    vec![
3403                        Encoding::PLAIN,
3404                        Encoding::DELTA_BYTE_ARRAY,
3405                        Encoding::DELTA_LENGTH_BYTE_ARRAY,
3406                    ]
3407                }
3408                DataType::Int64
3409                | DataType::Int32
3410                | DataType::Int16
3411                | DataType::Int8
3412                | DataType::UInt64
3413                | DataType::UInt32
3414                | DataType::UInt16
3415                | DataType::UInt8 => vec![
3416                    Encoding::PLAIN,
3417                    Encoding::DELTA_BINARY_PACKED,
3418                    Encoding::BYTE_STREAM_SPLIT,
3419                ],
3420                DataType::Float32 | DataType::Float64 => {
3421                    vec![Encoding::PLAIN, Encoding::BYTE_STREAM_SPLIT, Encoding::ALP]
3422                }
3423                _ => vec![Encoding::PLAIN],
3424            };
3425
3426            let expected_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
3427
3428            let row_group_sizes = [1024, SMALL_SIZE, SMALL_SIZE / 2, SMALL_SIZE / 2 + 1, 10];
3429
3430            let mut files = vec![];
3431            for dictionary_size in [0, 1, 1024] {
3432                for encoding in &encodings {
3433                    for version in [WriterVersion::PARQUET_1_0, WriterVersion::PARQUET_2_0] {
3434                        for row_group_size in row_group_sizes {
3435                            let mut builder = WriterProperties::builder()
3436                                .set_writer_version(version)
3437                                .set_max_row_group_row_count(Some(row_group_size))
3438                                .set_dictionary_enabled(dictionary_size != 0)
3439                                .set_dictionary_page_size_limit(dictionary_size.max(1))
3440                                .set_encoding(*encoding)
3441                                .set_bloom_filter_enabled(bloom_filter)
3442                                .set_bloom_filter_position(bloom_filter_position);
3443                            if let Some(ndv) = bloom_filter_ndv {
3444                                builder = builder.set_bloom_filter_max_ndv(ndv);
3445                            }
3446                            let props = builder.build();
3447
3448                            files.push(roundtrip_opts(&expected_batch, props))
3449                        }
3450                    }
3451                }
3452            }
3453            files
3454        }
3455    }
3456
3457    fn values_required<A, I>(iter: I) -> Vec<Bytes>
3458    where
3459        A: From<Vec<I::Item>> + Array + 'static,
3460        I: IntoIterator,
3461    {
3462        let raw_values: Vec<_> = iter.into_iter().collect();
3463        let values = Arc::new(A::from(raw_values));
3464        RoundTripTest::new(values).with_nullable(false).run()
3465    }
3466
3467    fn values_optional<A, I>(iter: I) -> Vec<Bytes>
3468    where
3469        A: From<Vec<Option<I::Item>>> + Array + 'static,
3470        I: IntoIterator,
3471    {
3472        let optional_raw_values: Vec<_> = iter
3473            .into_iter()
3474            .enumerate()
3475            .map(|(i, v)| if i % 2 == 0 { None } else { Some(v) })
3476            .collect();
3477        let optional_values = Arc::new(A::from(optional_raw_values));
3478        RoundTripTest::new(optional_values).run()
3479    }
3480
3481    fn required_and_optional<A, I>(iter: I)
3482    where
3483        A: From<Vec<I::Item>> + From<Vec<Option<I::Item>>> + Array + 'static,
3484        I: IntoIterator + Clone,
3485    {
3486        values_required::<A, I>(iter.clone());
3487        values_optional::<A, I>(iter);
3488    }
3489
3490    fn check_bloom_filter<T: AsBytes>(
3491        files: Vec<Bytes>,
3492        file_column: String,
3493        positive_values: Vec<T>,
3494        negative_values: Vec<T>,
3495    ) {
3496        files.into_iter().take(1).for_each(|file| {
3497            let file_reader = SerializedFileReader::new_with_options(
3498                file,
3499                ReadOptionsBuilder::new()
3500                    .with_reader_properties(
3501                        ReaderProperties::builder()
3502                            .set_read_bloom_filter(true)
3503                            .build(),
3504                    )
3505                    .build(),
3506            )
3507            .expect("Unable to open file as Parquet");
3508            let metadata = file_reader.metadata();
3509
3510            // Gets bloom filters from all row groups.
3511            let mut bloom_filters: Vec<_> = vec![];
3512            for (ri, row_group) in metadata.row_groups().iter().enumerate() {
3513                if let Some((column_index, _)) = row_group
3514                    .columns()
3515                    .iter()
3516                    .enumerate()
3517                    .find(|(_, column)| column.column_path().string() == file_column)
3518                {
3519                    let row_group_reader = file_reader
3520                        .get_row_group(ri)
3521                        .expect("Unable to read row group");
3522                    if let Some(sbbf) = row_group_reader.get_column_bloom_filter(column_index) {
3523                        bloom_filters.push(sbbf.clone());
3524                    } else {
3525                        panic!("No bloom filter for column named {file_column} found");
3526                    }
3527                } else {
3528                    panic!("No column named {file_column} found");
3529                }
3530            }
3531
3532            positive_values.iter().for_each(|value| {
3533                let found = bloom_filters.iter().find(|sbbf| sbbf.check(value));
3534                assert!(
3535                    found.is_some(),
3536                    "{}",
3537                    format!("Value {:?} should be in bloom filter", value.as_bytes())
3538                );
3539            });
3540
3541            negative_values.iter().for_each(|value| {
3542                let found = bloom_filters.iter().find(|sbbf| sbbf.check(value));
3543                assert!(
3544                    found.is_none(),
3545                    "{}",
3546                    format!("Value {:?} should not be in bloom filter", value.as_bytes())
3547                );
3548            });
3549        });
3550    }
3551
3552    #[test]
3553    #[cfg_attr(miri, ignore)] // Takes too long
3554    fn all_null_primitive_single_column() {
3555        let values = Arc::new(Int32Array::from(vec![None; SMALL_SIZE]));
3556        RoundTripTest::new(values).run();
3557    }
3558    #[test]
3559    #[cfg_attr(miri, ignore)] // Takes too long
3560    fn null_single_column() {
3561        let values = Arc::new(NullArray::new(SMALL_SIZE));
3562        RoundTripTest::new(values).run();
3563        // null arrays are always nullable, a test with non-nullable nulls fails
3564    }
3565
3566    #[test]
3567    #[cfg_attr(miri, ignore)] // Takes too long
3568    fn bool_single_column() {
3569        required_and_optional::<BooleanArray, _>(
3570            [true, false].iter().cycle().copied().take(SMALL_SIZE),
3571        );
3572    }
3573
3574    #[test]
3575    #[cfg_attr(miri, ignore)] // Takes too long
3576    fn bool_large_single_column() {
3577        let values = Arc::new(
3578            [None, Some(true), Some(false)]
3579                .iter()
3580                .cycle()
3581                .copied()
3582                .take(200_000)
3583                .collect::<BooleanArray>(),
3584        );
3585        let schema = Schema::new(vec![Field::new("col", values.data_type().clone(), true)]);
3586        let expected_batch = RecordBatch::try_new(Arc::new(schema), vec![values]).unwrap();
3587        let file = tempfile::tempfile().unwrap();
3588
3589        let mut writer =
3590            ArrowWriter::try_new(file.try_clone().unwrap(), expected_batch.schema(), None)
3591                .expect("Unable to write file");
3592        writer.write(&expected_batch).unwrap();
3593        writer.close().unwrap();
3594    }
3595
3596    #[test]
3597    fn check_page_offset_index_with_nan() {
3598        let values = Arc::new(Float64Array::from(vec![f64::NAN; 10]));
3599        let schema = Schema::new(vec![Field::new("col", DataType::Float64, true)]);
3600        let batch = RecordBatch::try_new(Arc::new(schema), vec![values]).unwrap();
3601
3602        let mut out = Vec::with_capacity(1024);
3603        let mut writer =
3604            ArrowWriter::try_new(&mut out, batch.schema(), None).expect("Unable to write file");
3605        writer.write(&batch).unwrap();
3606        let file_meta_data = writer.close().unwrap();
3607        for row_group in file_meta_data.row_groups() {
3608            for column in row_group.columns() {
3609                assert!(column.offset_index_offset().is_some());
3610                assert!(column.offset_index_length().is_some());
3611                assert!(column.column_index_offset().is_some());
3612                assert!(column.column_index_length().is_some());
3613            }
3614        }
3615        if let Some(page_index) = file_meta_data.page_index() {
3616            for rg in 0..file_meta_data.num_row_groups() {
3617                for col in 0..file_meta_data.row_group(rg).num_columns() {
3618                    let idx = page_index
3619                        .column_index(rg, col)
3620                        .expect("column index should exist");
3621                    assert!(idx.nan_counts().is_some());
3622                    let ColumnIndexMetaData::DOUBLE(float_idx) = idx else {
3623                        panic!("expected double statistics")
3624                    };
3625                    for i in 0..idx.num_pages() as usize {
3626                        assert_eq!(float_idx.nan_count(i), Some(10));
3627                        assert_eq!(
3628                            f64::NAN.total_cmp(float_idx.min_value(i).unwrap()),
3629                            Ordering::Equal
3630                        );
3631                        assert_eq!(
3632                            f64::NAN.total_cmp(float_idx.max_value(i).unwrap()),
3633                            Ordering::Equal
3634                        );
3635                    }
3636                }
3637            }
3638        } else {
3639            panic!("page index should be present");
3640        }
3641    }
3642
3643    #[test]
3644    fn check_page_offset_index_with_mixed_nan() {
3645        let schema = Arc::new(Schema::new(vec![Field::new(
3646            "col",
3647            DataType::Float64,
3648            true,
3649        )]));
3650
3651        let mut out = Vec::with_capacity(1024);
3652        let props = WriterProperties::builder()
3653            .set_data_page_row_count_limit(10)
3654            .build();
3655        let mut writer = ArrowWriter::try_new(&mut out, schema.clone(), Some(props))
3656            .expect("Unable to write file");
3657
3658        // write a page of all NaN (since batch min and max are NaN, global min/max are NaN)
3659        let values = Arc::new(Float64Array::from(vec![f64::NAN; 10]));
3660        let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
3661        writer.write(&batch).unwrap();
3662
3663        // write a page of all -NaN (batch min/max is -NaN, should update global min to -NaN)
3664        let values = Arc::new(Float64Array::from(vec![-f64::NAN; 10]));
3665        let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
3666        writer.write(&batch).unwrap();
3667
3668        // write a page of all 0 (non-NaN should override global min/max, now 0/0)
3669        let values = Arc::new(Float64Array::from(vec![0_f64; 10]));
3670        let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
3671        writer.write(&batch).unwrap();
3672
3673        // write a mixed page (should now have min -1, max 1)
3674        let values = Arc::new(Float64Array::from(vec![
3675            -1.0,
3676            0.0,
3677            f64::NAN,
3678            -f64::NAN,
3679            1.0,
3680        ]));
3681        let batch = RecordBatch::try_new(schema.clone(), vec![values]).unwrap();
3682        writer.write(&batch).unwrap();
3683
3684        let file_meta_data = writer.close().unwrap();
3685
3686        // check the column chunk stats are correct
3687        let col_stats = file_meta_data
3688            .row_group(0)
3689            .column(0)
3690            .statistics()
3691            .expect("missing column chunk statistics");
3692
3693        assert_eq!(col_stats.nan_count_opt(), Some(22));
3694        assert_eq!(col_stats.min_bytes_opt(), Some((-1.0f64).as_bytes()));
3695        assert_eq!(col_stats.max_bytes_opt(), Some(1.0f64.as_bytes()));
3696
3697        assert!(file_meta_data.page_index().is_some());
3698        let col_idx = &file_meta_data.page_index().unwrap().column_index(0, 0);
3699        assert_eq!(col_idx.as_ref().unwrap().num_pages(), 4);
3700
3701        // test each page
3702        let Some(ColumnIndexMetaData::DOUBLE(float_idx)) = col_idx else {
3703            panic!("expected double statistics")
3704        };
3705
3706        assert_eq!(float_idx.nan_counts, Some(vec![10, 10, 0, 2]));
3707        assert_eq!(
3708            f64::NAN.total_cmp(float_idx.min_value(0).unwrap()),
3709            Ordering::Equal
3710        );
3711        assert_eq!(
3712            f64::NAN.total_cmp(float_idx.max_value(0).unwrap()),
3713            Ordering::Equal
3714        );
3715        assert_eq!(
3716            (-f64::NAN).total_cmp(float_idx.min_value(1).unwrap()),
3717            Ordering::Equal
3718        );
3719        assert_eq!(
3720            (-f64::NAN).total_cmp(float_idx.max_value(1).unwrap()),
3721            Ordering::Equal
3722        );
3723        assert_eq!(float_idx.min_value(2), Some(&0.0));
3724        assert_eq!(float_idx.max_value(2), Some(&0.0));
3725        assert_eq!(float_idx.min_value(3), Some(&-1.0));
3726        assert_eq!(float_idx.max_value(3), Some(&1.0));
3727    }
3728
3729    #[test]
3730    #[cfg_attr(miri, ignore)] // Takes too long
3731    fn i8_single_column() {
3732        required_and_optional::<Int8Array, _>(0..SMALL_SIZE as i8);
3733    }
3734
3735    #[test]
3736    #[cfg_attr(miri, ignore)] // Takes too long
3737    fn i16_single_column() {
3738        required_and_optional::<Int16Array, _>(0..SMALL_SIZE as i16);
3739    }
3740
3741    #[test]
3742    #[cfg_attr(miri, ignore)] // Takes too long
3743    fn i32_single_column() {
3744        required_and_optional::<Int32Array, _>(0..SMALL_SIZE as i32);
3745    }
3746
3747    #[test]
3748    #[cfg_attr(miri, ignore)] // Takes too long
3749    fn i64_single_column() {
3750        required_and_optional::<Int64Array, _>(0..SMALL_SIZE as i64);
3751    }
3752
3753    #[test]
3754    #[cfg_attr(miri, ignore)] // Takes too long
3755    fn u8_single_column() {
3756        required_and_optional::<UInt8Array, _>(0..SMALL_SIZE as u8);
3757    }
3758
3759    #[test]
3760    #[cfg_attr(miri, ignore)] // Takes too long
3761    fn u16_single_column() {
3762        required_and_optional::<UInt16Array, _>(0..SMALL_SIZE as u16);
3763    }
3764
3765    #[test]
3766    #[cfg_attr(miri, ignore)] // Takes too long
3767    fn u32_single_column() {
3768        required_and_optional::<UInt32Array, _>(0..SMALL_SIZE as u32);
3769    }
3770
3771    #[test]
3772    #[cfg_attr(miri, ignore)] // Takes too long
3773    fn u64_single_column() {
3774        required_and_optional::<UInt64Array, _>(0..SMALL_SIZE as u64);
3775    }
3776
3777    #[test]
3778    #[cfg_attr(miri, ignore)] // Takes too long
3779    fn f32_single_column() {
3780        required_and_optional::<Float32Array, _>((0..SMALL_SIZE).map(|i| i as f32));
3781    }
3782
3783    #[test]
3784    #[cfg_attr(miri, ignore)] // Takes too long
3785    fn f64_single_column() {
3786        required_and_optional::<Float64Array, _>((0..SMALL_SIZE).map(|i| i as f64));
3787    }
3788
3789    // The timestamp array types don't implement From<Vec<T>> because they need the timezone
3790    // argument, and they also doesn't support building from a Vec<Option<T>>, so call
3791    // RoundTripTest manually instead of calling required_and_optional for these tests.
3792
3793    #[test]
3794    #[cfg_attr(miri, ignore)] // Takes too long
3795    fn timestamp_second_single_column() {
3796        let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
3797        let values = Arc::new(TimestampSecondArray::from(raw_values));
3798
3799        RoundTripTest::new(values).with_nullable(false).run();
3800    }
3801
3802    #[test]
3803    #[cfg_attr(miri, ignore)] // Takes too long
3804    fn timestamp_millisecond_single_column() {
3805        let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
3806        let values = Arc::new(TimestampMillisecondArray::from(raw_values));
3807
3808        RoundTripTest::new(values).with_nullable(false).run();
3809    }
3810
3811    #[test]
3812    #[cfg_attr(miri, ignore)] // Takes too long
3813    fn timestamp_microsecond_single_column() {
3814        let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
3815        let values = Arc::new(TimestampMicrosecondArray::from(raw_values));
3816
3817        RoundTripTest::new(values).with_nullable(false).run();
3818    }
3819
3820    #[test]
3821    #[cfg_attr(miri, ignore)] // Takes too long
3822    fn timestamp_nanosecond_single_column() {
3823        let raw_values: Vec<_> = (0..SMALL_SIZE as i64).collect();
3824        let values = Arc::new(TimestampNanosecondArray::from(raw_values));
3825
3826        RoundTripTest::new(values).with_nullable(false).run();
3827    }
3828
3829    #[test]
3830    #[cfg_attr(miri, ignore)] // Takes too long
3831    fn date32_single_column() {
3832        required_and_optional::<Date32Array, _>(0..SMALL_SIZE as i32);
3833    }
3834
3835    #[test]
3836    #[cfg_attr(miri, ignore)] // Takes too long
3837    fn date64_single_column() {
3838        // Date64 must be a multiple of 86400000, see ARROW-10925
3839        required_and_optional::<Date64Array, _>(
3840            (0..(SMALL_SIZE as i64 * 86400000)).step_by(86400000),
3841        );
3842    }
3843
3844    #[test]
3845    #[cfg_attr(miri, ignore)] // Takes too long
3846    fn time32_second_single_column() {
3847        required_and_optional::<Time32SecondArray, _>(0..SMALL_SIZE as i32);
3848    }
3849
3850    #[test]
3851    #[cfg_attr(miri, ignore)] // Takes too long
3852    fn time32_millisecond_single_column() {
3853        required_and_optional::<Time32MillisecondArray, _>(0..SMALL_SIZE as i32);
3854    }
3855
3856    #[test]
3857    #[cfg_attr(miri, ignore)] // Takes too long
3858    fn time64_microsecond_single_column() {
3859        required_and_optional::<Time64MicrosecondArray, _>(0..SMALL_SIZE as i64);
3860    }
3861
3862    #[test]
3863    #[cfg_attr(miri, ignore)] // Takes too long
3864    fn time64_nanosecond_single_column() {
3865        required_and_optional::<Time64NanosecondArray, _>(0..SMALL_SIZE as i64);
3866    }
3867
3868    #[test]
3869    #[cfg_attr(miri, ignore)] // Takes too long
3870    fn duration_second_single_column() {
3871        required_and_optional::<DurationSecondArray, _>(0..SMALL_SIZE as i64);
3872    }
3873
3874    #[test]
3875    #[cfg_attr(miri, ignore)] // Takes too long
3876    fn duration_millisecond_single_column() {
3877        required_and_optional::<DurationMillisecondArray, _>(0..SMALL_SIZE as i64);
3878    }
3879
3880    #[test]
3881    #[cfg_attr(miri, ignore)] // Takes too long
3882    fn duration_microsecond_single_column() {
3883        required_and_optional::<DurationMicrosecondArray, _>(0..SMALL_SIZE as i64);
3884    }
3885
3886    #[test]
3887    #[cfg_attr(miri, ignore)] // Takes too long
3888    fn duration_nanosecond_single_column() {
3889        required_and_optional::<DurationNanosecondArray, _>(0..SMALL_SIZE as i64);
3890    }
3891
3892    #[test]
3893    #[cfg_attr(miri, ignore)] // Takes too long
3894    fn interval_year_month_single_column() {
3895        required_and_optional::<IntervalYearMonthArray, _>(0..SMALL_SIZE as i32);
3896    }
3897
3898    #[test]
3899    #[cfg_attr(miri, ignore)] // Takes too long
3900    fn interval_day_time_single_column() {
3901        required_and_optional::<IntervalDayTimeArray, _>(vec![
3902            IntervalDayTime::new(0, 1),
3903            IntervalDayTime::new(0, 3),
3904            IntervalDayTime::new(3, -2),
3905            IntervalDayTime::new(-200, 4),
3906        ]);
3907    }
3908
3909    #[test]
3910    #[should_panic(
3911        expected = "Attempting to write an Arrow interval type MonthDayNano to parquet that is not yet implemented"
3912    )]
3913    fn interval_month_day_nano_single_column() {
3914        required_and_optional::<IntervalMonthDayNanoArray, _>(vec![
3915            IntervalMonthDayNano::new(0, 1, 5),
3916            IntervalMonthDayNano::new(0, 3, 2),
3917            IntervalMonthDayNano::new(3, -2, -5),
3918            IntervalMonthDayNano::new(-200, 4, -1),
3919        ]);
3920    }
3921
3922    #[test]
3923    #[cfg_attr(miri, ignore)] // Takes too long
3924    fn binary_single_column() {
3925        let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
3926        let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
3927        let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
3928
3929        // BinaryArrays can't be built from Vec<Option<&str>>, so only call `values_required`
3930        values_required::<BinaryArray, _>(many_vecs_iter);
3931    }
3932
3933    #[test]
3934    #[cfg_attr(miri, ignore)] // Takes too long
3935    fn binary_view_single_column() {
3936        let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
3937        let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
3938        let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
3939
3940        // BinaryArrays can't be built from Vec<Option<&str>>, so only call `values_required`
3941        values_required::<BinaryViewArray, _>(many_vecs_iter);
3942    }
3943
3944    #[test]
3945    #[cfg_attr(miri, ignore)] // Takes too long
3946    fn i32_column_bloom_filter_at_end() {
3947        let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
3948        let files = RoundTripTest::new(array)
3949            .with_nullable(false)
3950            .with_bloom_filter(true)
3951            .with_bloom_filter_position(BloomFilterPosition::End)
3952            .run();
3953
3954        check_bloom_filter(
3955            files,
3956            "col".to_string(),
3957            (0..SMALL_SIZE as i32).collect(),
3958            (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
3959        );
3960    }
3961
3962    #[test]
3963    #[cfg_attr(miri, ignore)] // Takes too long
3964    fn i32_column_bloom_filter() {
3965        let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
3966        let files = RoundTripTest::new(array)
3967            .with_nullable(false)
3968            .with_bloom_filter(true)
3969            .run();
3970
3971        check_bloom_filter(
3972            files,
3973            "col".to_string(),
3974            (0..SMALL_SIZE as i32).collect(),
3975            (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
3976        );
3977    }
3978
3979    fn write_with_bloom_filter(array: ArrayRef, dictionary_page_size_limit: usize) -> Bytes {
3980        let schema = Arc::new(Schema::new(vec![Field::new(
3981            "col",
3982            array.data_type().clone(),
3983            false,
3984        )]));
3985        let batch = RecordBatch::try_new(schema.clone(), vec![array]).unwrap();
3986        let props = WriterProperties::builder()
3987            .set_dictionary_enabled(true)
3988            .set_dictionary_page_size_limit(dictionary_page_size_limit)
3989            .set_write_batch_size(256)
3990            .set_bloom_filter_enabled(true)
3991            .build();
3992        let mut buf = Vec::new();
3993        let mut writer = ArrowWriter::try_new(&mut buf, schema, Some(props)).unwrap();
3994        writer.write(&batch).unwrap();
3995        writer.close().unwrap();
3996        Bytes::from(buf)
3997    }
3998
3999    fn data_page_encoding_mask(file: &Bytes) -> EncodingMask {
4000        let metadata = ParquetMetaDataReader::new().parse_and_finish(file).unwrap();
4001        *metadata
4002            .row_group(0)
4003            .column(0)
4004            .page_encoding_stats_mask()
4005            .unwrap()
4006    }
4007
4008    /// While a column is dictionary encoded the bloom filter is populated from the dictionary
4009    /// when it is flushed, so a chunk that stays dictionary encoded must still contain every value.
4010    #[test]
4011    fn string_column_bloom_filter_populated_from_dictionary() {
4012        let values: Vec<String> = (0..2000).map(|i| format!("value-{}", i % 10)).collect();
4013        let array = Arc::new(StringArray::from_iter_values(&values));
4014        let file = write_with_bloom_filter(array, 1024 * 1024);
4015        assert!(data_page_encoding_mask(&file).is_only(Encoding::RLE_DICTIONARY));
4016
4017        check_bloom_filter(
4018            vec![file],
4019            "col".to_string(),
4020            (0..10).map(|i| format!("value-{i}").into_bytes()).collect(),
4021            (10..20)
4022                .map(|i| format!("value-{i}").into_bytes())
4023                .collect(),
4024        );
4025    }
4026
4027    /// After falling back from dictionary encoding the filter holds the dictionary's values
4028    /// and every value written plain afterwards.
4029    #[test]
4030    fn string_column_bloom_filter_across_dictionary_fallback() {
4031        let values: Vec<String> = (0..2000).map(|i| format!("value-{i}")).collect();
4032        let array = Arc::new(StringArray::from_iter_values(&values));
4033        let file = write_with_bloom_filter(array, 1024);
4034        let encodings = data_page_encoding_mask(&file);
4035        assert!(
4036            encodings.is_set(Encoding::RLE_DICTIONARY) && encodings.is_set(Encoding::PLAIN),
4037            "expected dictionary and plain data pages, got {encodings:?}"
4038        );
4039
4040        check_bloom_filter(
4041            vec![file],
4042            "col".to_string(),
4043            values.into_iter().map(String::into_bytes).collect(),
4044            (2000..2010)
4045                .map(|i| format!("value-{i}").into_bytes())
4046                .collect(),
4047        );
4048    }
4049
4050    #[test]
4051    fn i64_column_bloom_filter_populated_from_dictionary() {
4052        let array = Arc::new(Int64Array::from_iter_values((0..2000).map(|i| i % 10)));
4053        let file = write_with_bloom_filter(array, 1024 * 1024);
4054        assert!(data_page_encoding_mask(&file).is_only(Encoding::RLE_DICTIONARY));
4055
4056        check_bloom_filter(
4057            vec![file],
4058            "col".to_string(),
4059            (0..10i64).collect(),
4060            (10..20i64).collect(),
4061        );
4062    }
4063
4064    #[test]
4065    fn i64_column_bloom_filter_across_dictionary_fallback() {
4066        let array = Arc::new(Int64Array::from_iter_values(0..2000i64));
4067        let file = write_with_bloom_filter(array, 1024);
4068        let encodings = data_page_encoding_mask(&file);
4069        assert!(
4070            encodings.is_set(Encoding::RLE_DICTIONARY) && encodings.is_set(Encoding::PLAIN),
4071            "expected dictionary and plain data pages, got {encodings:?}"
4072        );
4073
4074        check_bloom_filter(
4075            vec![file],
4076            "col".to_string(),
4077            (0..2000i64).collect(),
4078            (2000..2010i64).collect(),
4079        );
4080    }
4081
4082    /// Test that bloom filter folding produces correct results even when
4083    /// the configured NDV differs significantly from actual NDV.
4084    /// A large NDV means a larger initial filter that gets folded down;
4085    /// a small NDV means a smaller initial filter.
4086    #[test]
4087    #[cfg_attr(miri, ignore)] // Takes too long
4088    fn i32_column_bloom_filter_fixed_ndv() {
4089        let array = Arc::new(Int32Array::from_iter(0..SMALL_SIZE as i32));
4090
4091        // NDV much larger than actual distinct values — tests folding a large filter down
4092        let files = RoundTripTest::new(array.clone())
4093            .with_nullable(false)
4094            .with_bloom_filter(true)
4095            .with_bloom_filter_ndv(1_000_000)
4096            .run();
4097
4098        check_bloom_filter(
4099            files,
4100            "col".to_string(),
4101            (0..SMALL_SIZE as i32).collect(),
4102            (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
4103        );
4104
4105        // NDV smaller than actual distinct values — tests the underestimate path
4106        let files = RoundTripTest::new(array)
4107            .with_nullable(false)
4108            .with_bloom_filter(true)
4109            .with_bloom_filter_ndv(3)
4110            .run();
4111
4112        check_bloom_filter(
4113            files,
4114            "col".to_string(),
4115            (0..SMALL_SIZE as i32).collect(),
4116            (SMALL_SIZE as i32 + 1..SMALL_SIZE as i32 + 10).collect(),
4117        );
4118    }
4119
4120    #[test]
4121    #[cfg_attr(miri, ignore)] // Takes too long
4122    fn binary_column_bloom_filter() {
4123        let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
4124        let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
4125        let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
4126
4127        let array = Arc::new(BinaryArray::from_iter_values(many_vecs_iter));
4128        let files = RoundTripTest::new(array)
4129            .with_nullable(false)
4130            .with_bloom_filter(true)
4131            .run();
4132
4133        check_bloom_filter(
4134            files,
4135            "col".to_string(),
4136            many_vecs,
4137            vec![vec![(SMALL_SIZE + 1) as u8]],
4138        );
4139    }
4140
4141    #[test]
4142    #[cfg_attr(miri, ignore)] // Takes too long
4143    fn empty_string_null_column_bloom_filter() {
4144        let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
4145        let raw_strs = raw_values.iter().map(|s| s.as_str());
4146
4147        let array = Arc::new(StringArray::from_iter_values(raw_strs));
4148        let files = RoundTripTest::new(array)
4149            .with_nullable(false)
4150            .with_bloom_filter(true)
4151            .run();
4152
4153        let optional_raw_values: Vec<_> = raw_values
4154            .iter()
4155            .enumerate()
4156            .filter_map(|(i, v)| if i % 2 == 0 { None } else { Some(v.as_str()) })
4157            .collect();
4158        // For null slots, empty string should not be in bloom filter.
4159        check_bloom_filter(files, "col".to_string(), optional_raw_values, vec![""]);
4160    }
4161
4162    #[test]
4163    #[cfg_attr(miri, ignore)] // Takes too long
4164    fn large_binary_single_column() {
4165        let one_vec: Vec<u8> = (0..SMALL_SIZE as u8).collect();
4166        let many_vecs: Vec<_> = std::iter::repeat_n(one_vec, SMALL_SIZE).collect();
4167        let many_vecs_iter = many_vecs.iter().map(|v| v.as_slice());
4168
4169        // LargeBinaryArrays can't be built from Vec<Option<&str>>, so only call `values_required`
4170        values_required::<LargeBinaryArray, _>(many_vecs_iter);
4171    }
4172
4173    #[test]
4174    #[cfg_attr(miri, ignore)] // Takes too long
4175    fn fixed_size_binary_single_column() {
4176        let mut builder = FixedSizeBinaryBuilder::new(4);
4177        builder.append_value(b"0123").unwrap();
4178        builder.append_null();
4179        builder.append_value(b"8910").unwrap();
4180        builder.append_value(b"1112").unwrap();
4181        let array = Arc::new(builder.finish());
4182
4183        RoundTripTest::new(array).run();
4184    }
4185
4186    #[test]
4187    #[cfg_attr(miri, ignore)] // Takes too long
4188    fn string_single_column() {
4189        let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
4190        let raw_strs = raw_values.iter().map(|s| s.as_str());
4191
4192        required_and_optional::<StringArray, _>(raw_strs);
4193    }
4194
4195    #[test]
4196    #[cfg_attr(miri, ignore)] // Takes too long
4197    fn large_string_single_column() {
4198        let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
4199        let raw_strs = raw_values.iter().map(|s| s.as_str());
4200
4201        required_and_optional::<LargeStringArray, _>(raw_strs);
4202    }
4203
4204    #[test]
4205    #[cfg_attr(miri, ignore)] // Takes too long
4206    fn string_view_single_column() {
4207        let raw_values: Vec<_> = (0..SMALL_SIZE).map(|i| i.to_string()).collect();
4208        let raw_strs = raw_values.iter().map(|s| s.as_str());
4209
4210        required_and_optional::<StringViewArray, _>(raw_strs);
4211    }
4212
4213    #[test]
4214    fn null_list_single_column() {
4215        let null_field = Field::new_list_field(DataType::Null, true);
4216        let list_field = Field::new("emptylist", DataType::List(Arc::new(null_field)), true);
4217
4218        let schema = Schema::new(vec![list_field]);
4219
4220        // Build [[], null, [null, null]]
4221        let a_values = NullArray::new(2);
4222        let a_value_offsets = arrow::buffer::Buffer::from([0, 0, 0, 2].to_byte_slice());
4223        let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
4224            DataType::Null,
4225            true,
4226        ))))
4227        .len(3)
4228        .add_buffer(a_value_offsets)
4229        .null_bit_buffer(Some(Buffer::from([0b00000101])))
4230        .add_child_data(a_values.into_data())
4231        .build()
4232        .unwrap();
4233
4234        let a = ListArray::from(a_list_data);
4235
4236        assert!(a.is_valid(0));
4237        assert!(!a.is_valid(1));
4238        assert!(a.is_valid(2));
4239
4240        assert_eq!(a.value(0).len(), 0);
4241        assert_eq!(a.value(2).len(), 2);
4242        assert_eq!(a.value(2).logical_nulls().unwrap().null_count(), 2);
4243
4244        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
4245        roundtrip(batch, None);
4246    }
4247
4248    #[test]
4249    #[cfg_attr(miri, ignore)] // Takes too long
4250    fn list_single_column() {
4251        let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
4252        let a_value_offsets = arrow::buffer::Buffer::from([0, 1, 3, 3, 6, 10].to_byte_slice());
4253        let a_list_data = ArrayData::builder(DataType::List(Arc::new(Field::new_list_field(
4254            DataType::Int32,
4255            false,
4256        ))))
4257        .len(5)
4258        .add_buffer(a_value_offsets)
4259        .null_bit_buffer(Some(Buffer::from([0b00011011])))
4260        .add_child_data(a_values.into_data())
4261        .build()
4262        .unwrap();
4263
4264        assert_eq!(a_list_data.null_count(), 1);
4265
4266        let a = ListArray::from(a_list_data);
4267        let values = Arc::new(a);
4268
4269        RoundTripTest::new(values).run();
4270    }
4271
4272    #[test]
4273    #[cfg_attr(miri, ignore)] // Takes too long
4274    fn large_list_single_column() {
4275        let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
4276        let a_value_offsets = arrow::buffer::Buffer::from([0i64, 1, 3, 3, 6, 10].to_byte_slice());
4277        let a_list_data = ArrayData::builder(DataType::LargeList(Arc::new(Field::new(
4278            "large_item",
4279            DataType::Int32,
4280            true,
4281        ))))
4282        .len(5)
4283        .add_buffer(a_value_offsets)
4284        .add_child_data(a_values.into_data())
4285        .null_bit_buffer(Some(Buffer::from([0b00011011])))
4286        .build()
4287        .unwrap();
4288
4289        // I think this setup is incorrect because this should pass
4290        assert_eq!(a_list_data.null_count(), 1);
4291
4292        let a = LargeListArray::from(a_list_data);
4293        let values = Arc::new(a);
4294
4295        RoundTripTest::new(values).run();
4296    }
4297
4298    #[test]
4299    #[cfg_attr(miri, ignore)] // Takes too long
4300    fn list_nested_nulls() {
4301        use arrow::datatypes::Int32Type;
4302        let data = vec![
4303            Some(vec![Some(1)]),
4304            Some(vec![Some(2), Some(3)]),
4305            None,
4306            Some(vec![Some(4), Some(5), None]),
4307            Some(vec![None]),
4308            Some(vec![Some(6), Some(7)]),
4309        ];
4310
4311        let list = ListArray::from_iter_primitive::<Int32Type, _, _>(data.clone());
4312        RoundTripTest::new(Arc::new(list)).run();
4313
4314        let list = LargeListArray::from_iter_primitive::<Int32Type, _, _>(data);
4315        RoundTripTest::new(Arc::new(list)).run();
4316    }
4317
4318    #[test]
4319    #[cfg_attr(miri, ignore)] // Takes too long
4320    fn list_utf8_view_selective_padding_roundtrip() {
4321        let item = Arc::new(Field::new_list_field(DataType::Utf8View, true));
4322        let mut builder = ListBuilder::new(StringViewBuilder::new()).with_field(item);
4323        builder.values().append_value("a");
4324        builder.values().append_null();
4325        builder.append(true);
4326        // The null parent list covers selective padding dropping values below
4327        // the list definition level while preserving the preceding item null.
4328        builder.append(false);
4329        // The long string covers the non-inlined Utf8View buffer path.
4330        builder.values().append_value("large payload over 12 bytes");
4331        builder.append(true);
4332
4333        RoundTripTest::new(Arc::new(builder.finish())).run();
4334    }
4335
4336    #[test]
4337    #[cfg_attr(miri, ignore)] // Takes too long
4338    fn struct_single_column() {
4339        let a_values = Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]);
4340        let struct_field_a = Arc::new(Field::new("f", DataType::Int32, false));
4341        let s = StructArray::from(vec![(struct_field_a, Arc::new(a_values) as ArrayRef)]);
4342
4343        let values = Arc::new(s);
4344        RoundTripTest::new(values).with_nullable(false).run();
4345    }
4346
4347    #[test]
4348    fn list_and_map_coerced_names() {
4349        // Create map and list with non-Parquet naming
4350        let list_field =
4351            Field::new_list("my_list", Field::new("item", DataType::Int32, false), false);
4352        let map_field = Field::new_map(
4353            "my_map",
4354            "my_entries",
4355            Field::new("my_keys", DataType::Int32, false),
4356            Field::new("my_values", DataType::Int32, true),
4357            false,
4358            true,
4359        );
4360
4361        let list_array = create_random_array(&list_field, 100, 0.0, 0.0).unwrap();
4362        let map_array = create_random_array(&map_field, 100, 0.0, 0.0).unwrap();
4363
4364        let arrow_schema = Arc::new(Schema::new(vec![list_field, map_field]));
4365
4366        // Write data to Parquet but coerce names to match spec
4367        let props = Some(WriterProperties::builder().set_coerce_types(true).build());
4368        let file = tempfile::tempfile().unwrap();
4369        let mut writer =
4370            ArrowWriter::try_new(file.try_clone().unwrap(), arrow_schema.clone(), props).unwrap();
4371
4372        let batch = RecordBatch::try_new(arrow_schema, vec![list_array, map_array]).unwrap();
4373        writer.write(&batch).unwrap();
4374        let file_metadata = writer.close().unwrap();
4375
4376        let schema = file_metadata.file_metadata().schema();
4377        // Coerced name of "item" should be "element"
4378        let list_field = &schema.get_fields()[0].get_fields()[0];
4379        assert_eq!(list_field.get_fields()[0].name(), "element");
4380
4381        let map_field = &schema.get_fields()[1].get_fields()[0];
4382        // Coerced name of "entries" should be "key_value"
4383        assert_eq!(map_field.name(), "key_value");
4384        // Coerced name of "my_keys" should be "key"
4385        assert_eq!(map_field.get_fields()[0].name(), "key");
4386        // Coerced name of "my_values" should be "value"
4387        assert_eq!(map_field.get_fields()[1].name(), "value");
4388
4389        // Double check schema after reading from the file
4390        let reader = SerializedFileReader::new(file).unwrap();
4391        let file_schema = reader.metadata().file_metadata().schema();
4392        let fields = file_schema.get_fields();
4393        let list_field = &fields[0].get_fields()[0];
4394        assert_eq!(list_field.get_fields()[0].name(), "element");
4395        let map_field = &fields[1].get_fields()[0];
4396        assert_eq!(map_field.name(), "key_value");
4397        assert_eq!(map_field.get_fields()[0].name(), "key");
4398        assert_eq!(map_field.get_fields()[1].name(), "value");
4399    }
4400
4401    #[test]
4402    #[cfg_attr(miri, ignore)] // Takes too long
4403    fn fallback_flush_data_page() {
4404        //tests if the Fallback::flush_data_page clears all buffers correctly
4405        let raw_values: Vec<_> = (0..MEDIUM_SIZE).map(|i| i.to_string()).collect();
4406        let values = Arc::new(StringArray::from(raw_values));
4407        let encodings = vec![
4408            Encoding::DELTA_BYTE_ARRAY,
4409            Encoding::DELTA_LENGTH_BYTE_ARRAY,
4410        ];
4411        let data_type = values.data_type().clone();
4412        let schema = Arc::new(Schema::new(vec![Field::new("col", data_type, false)]));
4413        let expected_batch = RecordBatch::try_new(schema, vec![values]).unwrap();
4414
4415        let row_group_sizes = [1024, SMALL_SIZE, SMALL_SIZE / 2, SMALL_SIZE / 2 + 1, 10];
4416        let data_page_size_limit: usize = 32;
4417        let write_batch_size: usize = 16;
4418
4419        for encoding in &encodings {
4420            for row_group_size in row_group_sizes {
4421                let props = WriterProperties::builder()
4422                    .set_writer_version(WriterVersion::PARQUET_2_0)
4423                    .set_max_row_group_row_count(Some(row_group_size))
4424                    .set_dictionary_enabled(false)
4425                    .set_encoding(*encoding)
4426                    .set_data_page_size_limit(data_page_size_limit)
4427                    .set_write_batch_size(write_batch_size)
4428                    .build();
4429
4430                roundtrip_opts_with_array_validation(&expected_batch, props, |a, b| {
4431                    let string_array_a = StringArray::from(a.clone());
4432                    let string_array_b = StringArray::from(b.clone());
4433                    let vec_a: Vec<&str> = string_array_a.iter().map(|v| v.unwrap()).collect();
4434                    let vec_b: Vec<&str> = string_array_b.iter().map(|v| v.unwrap()).collect();
4435                    assert_eq!(
4436                        vec_a, vec_b,
4437                        "failed for encoder: {encoding:?} and row_group_size: {row_group_size:?}"
4438                    );
4439                });
4440            }
4441        }
4442    }
4443
4444    #[test]
4445    #[cfg_attr(miri, ignore)] // Takes too long
4446    fn arrow_writer_string_dictionary() {
4447        // define schema
4448        #[expect(deprecated)]
4449        let schema = Arc::new(Schema::new(vec![Field::new_dict(
4450            "dictionary",
4451            DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)),
4452            true,
4453            42,
4454            true,
4455        )]));
4456
4457        // create some data
4458        let d: Int32DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")]
4459            .iter()
4460            .copied()
4461            .collect();
4462
4463        // build a record batch
4464        RoundTripTest::new(Arc::new(d)).with_schema(schema).run();
4465    }
4466
4467    #[test]
4468    fn arrow_writer_test_type_compatibility() {
4469        fn ensure_compatible_write<T1, T2>(array1: T1, array2: T2, expected_result: T1)
4470        where
4471            T1: Array + 'static,
4472            T2: Array + 'static,
4473        {
4474            let schema1 = Arc::new(Schema::new(vec![Field::new(
4475                "a",
4476                array1.data_type().clone(),
4477                false,
4478            )]));
4479
4480            let file = tempfile().unwrap();
4481            let mut writer =
4482                ArrowWriter::try_new(file.try_clone().unwrap(), schema1.clone(), None).unwrap();
4483
4484            let rb1 = RecordBatch::try_new(schema1.clone(), vec![Arc::new(array1)]).unwrap();
4485            writer.write(&rb1).unwrap();
4486
4487            let schema2 = Arc::new(Schema::new(vec![Field::new(
4488                "a",
4489                array2.data_type().clone(),
4490                false,
4491            )]));
4492            let rb2 = RecordBatch::try_new(schema2, vec![Arc::new(array2)]).unwrap();
4493            writer.write(&rb2).unwrap();
4494
4495            writer.close().unwrap();
4496
4497            let mut record_batch_reader =
4498                ParquetRecordBatchReader::try_new(file.try_clone().unwrap(), 1024).unwrap();
4499            let actual_batch = record_batch_reader.next().unwrap().unwrap();
4500
4501            let expected_batch =
4502                RecordBatch::try_new(schema1, vec![Arc::new(expected_result)]).unwrap();
4503            assert_eq!(actual_batch, expected_batch);
4504        }
4505
4506        // check compatibility between native and dictionaries
4507
4508        ensure_compatible_write(
4509            DictionaryArray::new(
4510                UInt8Array::from_iter_values(vec![0]),
4511                Arc::new(StringArray::from_iter_values(vec!["parquet"])),
4512            ),
4513            StringArray::from_iter_values(vec!["barquet"]),
4514            DictionaryArray::new(
4515                UInt8Array::from_iter_values(vec![0, 1]),
4516                Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])),
4517            ),
4518        );
4519
4520        ensure_compatible_write(
4521            StringArray::from_iter_values(vec!["parquet"]),
4522            DictionaryArray::new(
4523                UInt8Array::from_iter_values(vec![0]),
4524                Arc::new(StringArray::from_iter_values(vec!["barquet"])),
4525            ),
4526            StringArray::from_iter_values(vec!["parquet", "barquet"]),
4527        );
4528
4529        // check compatibility between dictionaries with different key types
4530
4531        ensure_compatible_write(
4532            DictionaryArray::new(
4533                UInt8Array::from_iter_values(vec![0]),
4534                Arc::new(StringArray::from_iter_values(vec!["parquet"])),
4535            ),
4536            DictionaryArray::new(
4537                UInt16Array::from_iter_values(vec![0]),
4538                Arc::new(StringArray::from_iter_values(vec!["barquet"])),
4539            ),
4540            DictionaryArray::new(
4541                UInt8Array::from_iter_values(vec![0, 1]),
4542                Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])),
4543            ),
4544        );
4545
4546        // check compatibility between dictionaries with different value types
4547        ensure_compatible_write(
4548            DictionaryArray::new(
4549                UInt8Array::from_iter_values(vec![0]),
4550                Arc::new(StringArray::from_iter_values(vec!["parquet"])),
4551            ),
4552            DictionaryArray::new(
4553                UInt8Array::from_iter_values(vec![0]),
4554                Arc::new(LargeStringArray::from_iter_values(vec!["barquet"])),
4555            ),
4556            DictionaryArray::new(
4557                UInt8Array::from_iter_values(vec![0, 1]),
4558                Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])),
4559            ),
4560        );
4561
4562        // check compatibility between a dictionary and a native array with a different type
4563        ensure_compatible_write(
4564            DictionaryArray::new(
4565                UInt8Array::from_iter_values(vec![0]),
4566                Arc::new(StringArray::from_iter_values(vec!["parquet"])),
4567            ),
4568            LargeStringArray::from_iter_values(vec!["barquet"]),
4569            DictionaryArray::new(
4570                UInt8Array::from_iter_values(vec![0, 1]),
4571                Arc::new(StringArray::from_iter_values(vec!["parquet", "barquet"])),
4572            ),
4573        );
4574
4575        // check compatibility for string types
4576
4577        ensure_compatible_write(
4578            StringArray::from_iter_values(vec!["parquet"]),
4579            LargeStringArray::from_iter_values(vec!["barquet"]),
4580            StringArray::from_iter_values(vec!["parquet", "barquet"]),
4581        );
4582
4583        ensure_compatible_write(
4584            LargeStringArray::from_iter_values(vec!["parquet"]),
4585            StringArray::from_iter_values(vec!["barquet"]),
4586            LargeStringArray::from_iter_values(vec!["parquet", "barquet"]),
4587        );
4588
4589        ensure_compatible_write(
4590            StringArray::from_iter_values(vec!["parquet"]),
4591            StringViewArray::from_iter_values(vec!["barquet"]),
4592            StringArray::from_iter_values(vec!["parquet", "barquet"]),
4593        );
4594
4595        ensure_compatible_write(
4596            StringViewArray::from_iter_values(vec!["parquet"]),
4597            StringArray::from_iter_values(vec!["barquet"]),
4598            StringViewArray::from_iter_values(vec!["parquet", "barquet"]),
4599        );
4600
4601        ensure_compatible_write(
4602            LargeStringArray::from_iter_values(vec!["parquet"]),
4603            StringViewArray::from_iter_values(vec!["barquet"]),
4604            LargeStringArray::from_iter_values(vec!["parquet", "barquet"]),
4605        );
4606
4607        ensure_compatible_write(
4608            StringViewArray::from_iter_values(vec!["parquet"]),
4609            LargeStringArray::from_iter_values(vec!["barquet"]),
4610            StringViewArray::from_iter_values(vec!["parquet", "barquet"]),
4611        );
4612
4613        // check compatibility for binary types
4614
4615        ensure_compatible_write(
4616            BinaryArray::from_iter_values(vec![b"parquet"]),
4617            LargeBinaryArray::from_iter_values(vec![b"barquet"]),
4618            BinaryArray::from_iter_values(vec![b"parquet", b"barquet"]),
4619        );
4620
4621        ensure_compatible_write(
4622            LargeBinaryArray::from_iter_values(vec![b"parquet"]),
4623            BinaryArray::from_iter_values(vec![b"barquet"]),
4624            LargeBinaryArray::from_iter_values(vec![b"parquet", b"barquet"]),
4625        );
4626
4627        ensure_compatible_write(
4628            BinaryArray::from_iter_values(vec![b"parquet"]),
4629            BinaryViewArray::from_iter_values(vec![b"barquet"]),
4630            BinaryArray::from_iter_values(vec![b"parquet", b"barquet"]),
4631        );
4632
4633        ensure_compatible_write(
4634            BinaryViewArray::from_iter_values(vec![b"parquet"]),
4635            BinaryArray::from_iter_values(vec![b"barquet"]),
4636            BinaryViewArray::from_iter_values(vec![b"parquet", b"barquet"]),
4637        );
4638
4639        ensure_compatible_write(
4640            BinaryViewArray::from_iter_values(vec![b"parquet"]),
4641            LargeBinaryArray::from_iter_values(vec![b"barquet"]),
4642            BinaryViewArray::from_iter_values(vec![b"parquet", b"barquet"]),
4643        );
4644
4645        ensure_compatible_write(
4646            LargeBinaryArray::from_iter_values(vec![b"parquet"]),
4647            BinaryViewArray::from_iter_values(vec![b"barquet"]),
4648            LargeBinaryArray::from_iter_values(vec![b"parquet", b"barquet"]),
4649        );
4650
4651        // check compatibility for list types
4652
4653        let list_field_metadata = HashMap::from_iter(vec![(
4654            PARQUET_FIELD_ID_META_KEY.to_string(),
4655            "1".to_string(),
4656        )]);
4657        let list_field = Field::new_list_field(DataType::Int32, false);
4658
4659        let values1 = Arc::new(Int32Array::from(vec![0, 1, 2, 3, 4]));
4660        let offsets1 = OffsetBuffer::new(vec![0, 2, 5].into());
4661
4662        let values2 = Arc::new(Int32Array::from(vec![5, 6, 7, 8, 9]));
4663        let offsets2 = OffsetBuffer::new(vec![0, 3, 5].into());
4664
4665        let values_expected = Arc::new(Int32Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]));
4666        let offsets_expected = OffsetBuffer::new(vec![0, 2, 5, 8, 10].into());
4667
4668        ensure_compatible_write(
4669            // when the initial schema has the metadata ...
4670            ListArray::try_new(
4671                Arc::new(
4672                    list_field
4673                        .clone()
4674                        .with_metadata(list_field_metadata.clone()),
4675                ),
4676                offsets1,
4677                values1,
4678                None,
4679            )
4680            .unwrap(),
4681            // ... and some intermediate schema doesn't have the metadata
4682            ListArray::try_new(Arc::new(list_field.clone()), offsets2, values2, None).unwrap(),
4683            // ... the write will still go through, and the resulting schema will inherit the initial metadata
4684            ListArray::try_new(
4685                Arc::new(
4686                    list_field
4687                        .clone()
4688                        .with_metadata(list_field_metadata.clone()),
4689                ),
4690                offsets_expected,
4691                values_expected,
4692                None,
4693            )
4694            .unwrap(),
4695        );
4696    }
4697
4698    #[test]
4699    #[cfg_attr(miri, ignore)] // Takes too long
4700    fn arrow_writer_primitive_dictionary() {
4701        // define schema
4702        #[expect(deprecated)]
4703        let schema = Arc::new(Schema::new(vec![Field::new_dict(
4704            "dictionary",
4705            DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::UInt32)),
4706            true,
4707            42,
4708            true,
4709        )]));
4710
4711        // create some data
4712        let mut builder = PrimitiveDictionaryBuilder::<UInt8Type, UInt32Type>::new();
4713        builder.append(12345678).unwrap();
4714        builder.append_null();
4715        builder.append(22345678).unwrap();
4716        builder.append(12345678).unwrap();
4717        let d = builder.finish();
4718
4719        RoundTripTest::new(Arc::new(d)).with_schema(schema).run();
4720    }
4721
4722    #[test]
4723    #[cfg_attr(miri, ignore)] // Takes too long
4724    fn arrow_writer_decimal32_dictionary() {
4725        let integers = vec![12345, 56789, 34567];
4726
4727        let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]);
4728
4729        let values = Decimal32Array::from(integers.clone())
4730            .with_precision_and_scale(5, 2)
4731            .unwrap();
4732
4733        let array = DictionaryArray::new(keys, Arc::new(values));
4734        RoundTripTest::new(Arc::new(array.clone())).run();
4735
4736        let values = Decimal32Array::from(integers)
4737            .with_precision_and_scale(9, 2)
4738            .unwrap();
4739
4740        let array = array.with_values(Arc::new(values));
4741        RoundTripTest::new(Arc::new(array)).run();
4742    }
4743
4744    #[test]
4745    #[cfg_attr(miri, ignore)] // Takes too long
4746    fn arrow_writer_decimal64_dictionary() {
4747        let integers = vec![12345, 56789, 34567];
4748
4749        let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]);
4750
4751        let values = Decimal64Array::from(integers.clone())
4752            .with_precision_and_scale(5, 2)
4753            .unwrap();
4754
4755        let array = DictionaryArray::new(keys, Arc::new(values));
4756        RoundTripTest::new(Arc::new(array.clone())).run();
4757
4758        let values = Decimal64Array::from(integers)
4759            .with_precision_and_scale(12, 2)
4760            .unwrap();
4761
4762        let array = array.with_values(Arc::new(values));
4763        RoundTripTest::new(Arc::new(array)).run();
4764    }
4765
4766    #[test]
4767    #[cfg_attr(miri, ignore)] // Takes too long
4768    fn arrow_writer_decimal128_dictionary() {
4769        let integers = vec![12345, 56789, 34567];
4770
4771        let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]);
4772
4773        let values = Decimal128Array::from(integers.clone())
4774            .with_precision_and_scale(5, 2)
4775            .unwrap();
4776
4777        let array = DictionaryArray::new(keys, Arc::new(values));
4778        RoundTripTest::new(Arc::new(array.clone())).run();
4779
4780        let values = Decimal128Array::from(integers)
4781            .with_precision_and_scale(12, 2)
4782            .unwrap();
4783
4784        let array = array.with_values(Arc::new(values));
4785        RoundTripTest::new(Arc::new(array)).run();
4786    }
4787
4788    #[test]
4789    #[cfg_attr(miri, ignore)] // Takes too long
4790    fn arrow_writer_decimal256_dictionary() {
4791        let integers = vec![
4792            i256::from_i128(12345),
4793            i256::from_i128(56789),
4794            i256::from_i128(34567),
4795        ];
4796
4797        let keys = UInt8Array::from(vec![Some(0), None, Some(1), Some(2), Some(1)]);
4798
4799        let values = Decimal256Array::from(integers.clone())
4800            .with_precision_and_scale(5, 2)
4801            .unwrap();
4802
4803        let array = DictionaryArray::new(keys, Arc::new(values));
4804        RoundTripTest::new(Arc::new(array.clone())).run();
4805
4806        let values = Decimal256Array::from(integers)
4807            .with_precision_and_scale(12, 2)
4808            .unwrap();
4809
4810        let array = array.with_values(Arc::new(values));
4811        RoundTripTest::new(Arc::new(array)).run();
4812    }
4813
4814    #[test]
4815    #[cfg_attr(miri, ignore)] // Takes too long
4816    fn arrow_writer_string_dictionary_unsigned_index() {
4817        // define schema
4818        #[expect(deprecated)]
4819        let schema = Arc::new(Schema::new(vec![Field::new_dict(
4820            "dictionary",
4821            DataType::Dictionary(Box::new(DataType::UInt8), Box::new(DataType::Utf8)),
4822            true,
4823            42,
4824            true,
4825        )]));
4826
4827        // create some data
4828        let d: UInt8DictionaryArray = [Some("alpha"), None, Some("beta"), Some("alpha")]
4829            .iter()
4830            .copied()
4831            .collect();
4832
4833        RoundTripTest::new(Arc::new(d)).with_schema(schema).run();
4834    }
4835
4836    #[test]
4837    #[cfg_attr(miri, ignore)] // Takes too long
4838    fn u32_min_max() {
4839        // check values roundtrip through parquet
4840        let src = [
4841            u32::MIN,
4842            1,
4843            (i32::MAX as u32) - 1,
4844            i32::MAX as u32,
4845            (i32::MAX as u32) + 1,
4846            u32::MAX - 1,
4847            u32::MAX,
4848        ];
4849        let values = Arc::new(UInt32Array::from_iter_values(src.iter().copied()));
4850        let files = RoundTripTest::new(values).with_nullable(false).run();
4851
4852        for file in files {
4853            // check statistics are valid
4854            let reader = SerializedFileReader::new(file).unwrap();
4855            let metadata = reader.metadata();
4856
4857            let mut row_offset = 0;
4858            for row_group in metadata.row_groups() {
4859                assert_eq!(row_group.num_columns(), 1);
4860                let column = row_group.column(0);
4861
4862                let num_values = column.num_values() as usize;
4863                let src_slice = &src[row_offset..row_offset + num_values];
4864                row_offset += column.num_values() as usize;
4865
4866                let stats = column.statistics().unwrap();
4867                if let Statistics::Int32(stats) = stats {
4868                    assert_eq!(
4869                        *stats.min_opt().unwrap() as u32,
4870                        *src_slice.iter().min().unwrap()
4871                    );
4872                    assert_eq!(
4873                        *stats.max_opt().unwrap() as u32,
4874                        *src_slice.iter().max().unwrap()
4875                    );
4876                } else {
4877                    panic!("Statistics::Int32 missing")
4878                }
4879            }
4880        }
4881    }
4882
4883    #[test]
4884    #[cfg_attr(miri, ignore)] // Takes too long
4885    fn u64_min_max() {
4886        // check values roundtrip through parquet
4887        let src = [
4888            u64::MIN,
4889            1,
4890            (i64::MAX as u64) - 1,
4891            i64::MAX as u64,
4892            (i64::MAX as u64) + 1,
4893            u64::MAX - 1,
4894            u64::MAX,
4895        ];
4896        let values = Arc::new(UInt64Array::from_iter_values(src.iter().copied()));
4897        let files = RoundTripTest::new(values).with_nullable(false).run();
4898
4899        for file in files {
4900            // check statistics are valid
4901            let reader = SerializedFileReader::new(file).unwrap();
4902            let metadata = reader.metadata();
4903
4904            let mut row_offset = 0;
4905            for row_group in metadata.row_groups() {
4906                assert_eq!(row_group.num_columns(), 1);
4907                let column = row_group.column(0);
4908
4909                let num_values = column.num_values() as usize;
4910                let src_slice = &src[row_offset..row_offset + num_values];
4911                row_offset += column.num_values() as usize;
4912
4913                let stats = column.statistics().unwrap();
4914                if let Statistics::Int64(stats) = stats {
4915                    assert_eq!(
4916                        *stats.min_opt().unwrap() as u64,
4917                        *src_slice.iter().min().unwrap()
4918                    );
4919                    assert_eq!(
4920                        *stats.max_opt().unwrap() as u64,
4921                        *src_slice.iter().max().unwrap()
4922                    );
4923                } else {
4924                    panic!("Statistics::Int64 missing")
4925                }
4926            }
4927        }
4928    }
4929
4930    #[test]
4931    #[cfg_attr(miri, ignore)] // Takes too long
4932    fn statistics_null_counts_only_nulls() {
4933        // check that null-count statistics for "only NULL"-columns are correct
4934        let values = Arc::new(UInt64Array::from(vec![None, None]));
4935        let files = RoundTripTest::new(values).run();
4936
4937        for file in files {
4938            // check statistics are valid
4939            let reader = SerializedFileReader::new(file).unwrap();
4940            let metadata = reader.metadata();
4941            assert_eq!(metadata.num_row_groups(), 1);
4942            let row_group = metadata.row_group(0);
4943            assert_eq!(row_group.num_columns(), 1);
4944            let column = row_group.column(0);
4945            let stats = column.statistics().unwrap();
4946            assert_eq!(stats.null_count_opt(), Some(2));
4947        }
4948    }
4949
4950    #[test]
4951    #[cfg_attr(miri, ignore)] // Takes too long
4952    fn test_list_of_struct_roundtrip() {
4953        // define schema
4954        let int_field = Field::new("a", DataType::Int32, true);
4955        let int_field2 = Field::new("b", DataType::Int32, true);
4956
4957        let int_builder = Int32Builder::with_capacity(10);
4958        let int_builder2 = Int32Builder::with_capacity(10);
4959
4960        let struct_builder = StructBuilder::new(
4961            vec![int_field, int_field2],
4962            vec![Box::new(int_builder), Box::new(int_builder2)],
4963        );
4964        let mut list_builder = ListBuilder::new(struct_builder);
4965
4966        // Construct the following array
4967        // [{a: 1, b: 2}], [], null, [null, null], [{a: null, b: 3}], [{a: 2, b: null}]
4968
4969        // [{a: 1, b: 2}]
4970        let values = list_builder.values();
4971        values
4972            .field_builder::<Int32Builder>(0)
4973            .unwrap()
4974            .append_value(1);
4975        values
4976            .field_builder::<Int32Builder>(1)
4977            .unwrap()
4978            .append_value(2);
4979        values.append(true);
4980        list_builder.append(true);
4981
4982        // []
4983        list_builder.append(true);
4984
4985        // null
4986        list_builder.append(false);
4987
4988        // [null, null]
4989        let values = list_builder.values();
4990        values
4991            .field_builder::<Int32Builder>(0)
4992            .unwrap()
4993            .append_null();
4994        values
4995            .field_builder::<Int32Builder>(1)
4996            .unwrap()
4997            .append_null();
4998        values.append(false);
4999        values
5000            .field_builder::<Int32Builder>(0)
5001            .unwrap()
5002            .append_null();
5003        values
5004            .field_builder::<Int32Builder>(1)
5005            .unwrap()
5006            .append_null();
5007        values.append(false);
5008        list_builder.append(true);
5009
5010        // [{a: null, b: 3}]
5011        let values = list_builder.values();
5012        values
5013            .field_builder::<Int32Builder>(0)
5014            .unwrap()
5015            .append_null();
5016        values
5017            .field_builder::<Int32Builder>(1)
5018            .unwrap()
5019            .append_value(3);
5020        values.append(true);
5021        list_builder.append(true);
5022
5023        // [{a: 2, b: null}]
5024        let values = list_builder.values();
5025        values
5026            .field_builder::<Int32Builder>(0)
5027            .unwrap()
5028            .append_value(2);
5029        values
5030            .field_builder::<Int32Builder>(1)
5031            .unwrap()
5032            .append_null();
5033        values.append(true);
5034        list_builder.append(true);
5035
5036        let array = Arc::new(list_builder.finish());
5037
5038        RoundTripTest::new(array).run();
5039    }
5040
5041    fn row_group_sizes(metadata: &ParquetMetaData) -> Vec<i64> {
5042        metadata.row_groups().iter().map(|x| x.num_rows()).collect()
5043    }
5044
5045    #[test]
5046    fn test_aggregates_records() {
5047        let arrays = [
5048            Int32Array::from((0..100).collect::<Vec<_>>()),
5049            Int32Array::from((0..50).collect::<Vec<_>>()),
5050            Int32Array::from((200..500).collect::<Vec<_>>()),
5051        ];
5052
5053        let schema = Arc::new(Schema::new(vec![Field::new(
5054            "int",
5055            ArrowDataType::Int32,
5056            false,
5057        )]));
5058
5059        let file = tempfile::tempfile().unwrap();
5060
5061        let props = WriterProperties::builder()
5062            .set_max_row_group_row_count(Some(200))
5063            .build();
5064
5065        let mut writer =
5066            ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
5067
5068        for array in arrays {
5069            let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
5070            writer.write(&batch).unwrap();
5071        }
5072
5073        writer.close().unwrap();
5074
5075        let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
5076        assert_eq!(&row_group_sizes(builder.metadata()), &[200, 200, 50]);
5077
5078        let batches = builder
5079            .with_batch_size(100)
5080            .build()
5081            .unwrap()
5082            .collect::<ArrowResult<Vec<_>>>()
5083            .unwrap();
5084
5085        assert_eq!(batches.len(), 5);
5086        assert!(batches.iter().all(|x| x.num_columns() == 1));
5087
5088        let batch_sizes: Vec<_> = batches.iter().map(|x| x.num_rows()).collect();
5089
5090        assert_eq!(&batch_sizes, &[100, 100, 100, 100, 50]);
5091
5092        let values: Vec<_> = batches
5093            .iter()
5094            .flat_map(|x| {
5095                x.column(0)
5096                    .as_any()
5097                    .downcast_ref::<Int32Array>()
5098                    .unwrap()
5099                    .values()
5100                    .iter()
5101                    .copied()
5102            })
5103            .collect();
5104
5105        let expected_values: Vec<_> = [0..100, 0..50, 200..500].into_iter().flatten().collect();
5106        assert_eq!(&values, &expected_values)
5107    }
5108
5109    #[test]
5110    fn complex_aggregate() {
5111        // Tests aggregating nested data
5112        let field_a = Arc::new(Field::new("leaf_a", DataType::Int32, false));
5113        let field_b = Arc::new(Field::new("leaf_b", DataType::Int32, true));
5114        let struct_a = Arc::new(Field::new(
5115            "struct_a",
5116            DataType::Struct(vec![field_a.clone(), field_b.clone()].into()),
5117            true,
5118        ));
5119
5120        let list_a = Arc::new(Field::new("list", DataType::List(struct_a), true));
5121        let struct_b = Arc::new(Field::new(
5122            "struct_b",
5123            DataType::Struct(vec![list_a.clone()].into()),
5124            false,
5125        ));
5126
5127        let schema = Arc::new(Schema::new(vec![struct_b]));
5128
5129        // create nested data
5130        let field_a_array = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
5131        let field_b_array =
5132            Int32Array::from_iter(vec![Some(1), None, Some(2), None, None, Some(6)]);
5133
5134        let struct_a_array = StructArray::from(vec![
5135            (field_a.clone(), Arc::new(field_a_array) as ArrayRef),
5136            (field_b.clone(), Arc::new(field_b_array) as ArrayRef),
5137        ]);
5138
5139        let list_data = ArrayDataBuilder::new(list_a.data_type().clone())
5140            .len(5)
5141            .add_buffer(Buffer::from_iter(vec![
5142                0_i32, 1_i32, 1_i32, 3_i32, 3_i32, 5_i32,
5143            ]))
5144            .null_bit_buffer(Some(Buffer::from_iter(vec![
5145                true, false, true, false, true,
5146            ])))
5147            .child_data(vec![struct_a_array.into_data()])
5148            .build()
5149            .unwrap();
5150
5151        let list_a_array = Arc::new(ListArray::from(list_data)) as ArrayRef;
5152        let struct_b_array = StructArray::from(vec![(list_a.clone(), list_a_array)]);
5153
5154        let batch1 =
5155            RecordBatch::try_from_iter(vec![("struct_b", Arc::new(struct_b_array) as ArrayRef)])
5156                .unwrap();
5157
5158        let field_a_array = Int32Array::from(vec![6, 7, 8, 9, 10]);
5159        let field_b_array = Int32Array::from_iter(vec![None, None, None, Some(1), None]);
5160
5161        let struct_a_array = StructArray::from(vec![
5162            (field_a, Arc::new(field_a_array) as ArrayRef),
5163            (field_b, Arc::new(field_b_array) as ArrayRef),
5164        ]);
5165
5166        let list_data = ArrayDataBuilder::new(list_a.data_type().clone())
5167            .len(2)
5168            .add_buffer(Buffer::from_iter(vec![0_i32, 4_i32, 5_i32]))
5169            .child_data(vec![struct_a_array.into_data()])
5170            .build()
5171            .unwrap();
5172
5173        let list_a_array = Arc::new(ListArray::from(list_data)) as ArrayRef;
5174        let struct_b_array = StructArray::from(vec![(list_a, list_a_array)]);
5175
5176        let batch2 =
5177            RecordBatch::try_from_iter(vec![("struct_b", Arc::new(struct_b_array) as ArrayRef)])
5178                .unwrap();
5179
5180        let batches = &[batch1, batch2];
5181
5182        // Verify data is as expected
5183
5184        let expected = r"
5185            +-------------------------------------------------------------------------------------------------------+
5186            | struct_b                                                                                              |
5187            +-------------------------------------------------------------------------------------------------------+
5188            | {list: [{leaf_a: 1, leaf_b: 1}]}                                                                      |
5189            | {list: }                                                                                              |
5190            | {list: [{leaf_a: 2, leaf_b: }, {leaf_a: 3, leaf_b: 2}]}                                               |
5191            | {list: }                                                                                              |
5192            | {list: [{leaf_a: 4, leaf_b: }, {leaf_a: 5, leaf_b: }]}                                                |
5193            | {list: [{leaf_a: 6, leaf_b: }, {leaf_a: 7, leaf_b: }, {leaf_a: 8, leaf_b: }, {leaf_a: 9, leaf_b: 1}]} |
5194            | {list: [{leaf_a: 10, leaf_b: }]}                                                                      |
5195            +-------------------------------------------------------------------------------------------------------+
5196        ".trim().split('\n').map(|x| x.trim()).collect::<Vec<_>>().join("\n");
5197
5198        let actual = pretty_format_batches(batches).unwrap().to_string();
5199        assert_eq!(actual, expected);
5200
5201        // Write data
5202        let file = tempfile::tempfile().unwrap();
5203        let props = WriterProperties::builder()
5204            .set_max_row_group_row_count(Some(6))
5205            .build();
5206
5207        let mut writer =
5208            ArrowWriter::try_new(file.try_clone().unwrap(), schema, Some(props)).unwrap();
5209
5210        for batch in batches {
5211            writer.write(batch).unwrap();
5212        }
5213        writer.close().unwrap();
5214
5215        // Read Data
5216        // Should have written entire first batch and first row of second to the first row group
5217        // leaving a single row in the second row group
5218
5219        let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
5220        assert_eq!(&row_group_sizes(builder.metadata()), &[6, 1]);
5221
5222        let batches = builder
5223            .with_batch_size(2)
5224            .build()
5225            .unwrap()
5226            .collect::<ArrowResult<Vec<_>>>()
5227            .unwrap();
5228
5229        assert_eq!(batches.len(), 4);
5230        let batch_counts: Vec<_> = batches.iter().map(|x| x.num_rows()).collect();
5231        assert_eq!(&batch_counts, &[2, 2, 2, 1]);
5232
5233        let actual = pretty_format_batches(&batches).unwrap().to_string();
5234        assert_eq!(actual, expected);
5235    }
5236
5237    #[test]
5238    fn test_arrow_writer_metadata() {
5239        let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
5240        let file_schema = batch_schema.clone().with_metadata([("foo", "bar")]);
5241
5242        let batch = RecordBatch::try_new(
5243            Arc::new(batch_schema),
5244            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5245        )
5246        .unwrap();
5247
5248        let mut buf = Vec::with_capacity(1024);
5249        let mut writer = ArrowWriter::try_new(&mut buf, Arc::new(file_schema), None).unwrap();
5250        writer.write(&batch).unwrap();
5251        writer.close().unwrap();
5252    }
5253
5254    #[test]
5255    fn test_arrow_writer_nullable() {
5256        let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
5257        let file_schema = Schema::new(vec![Field::new("int32", DataType::Int32, true)]);
5258        let file_schema = Arc::new(file_schema);
5259
5260        let batch = RecordBatch::try_new(
5261            Arc::new(batch_schema),
5262            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5263        )
5264        .unwrap();
5265
5266        let mut buf = Vec::with_capacity(1024);
5267        let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), None).unwrap();
5268        writer.write(&batch).unwrap();
5269        writer.close().unwrap();
5270
5271        let mut read = ParquetRecordBatchReader::try_new(Bytes::from(buf), 1024).unwrap();
5272        let back = read.next().unwrap().unwrap();
5273        assert_eq!(back.schema(), file_schema);
5274        assert_ne!(back.schema(), batch.schema());
5275        assert_eq!(back.column(0).as_ref(), batch.column(0).as_ref());
5276    }
5277
5278    #[test]
5279    fn in_progress_accounting() {
5280        // define schema
5281        let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
5282
5283        // create some data
5284        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
5285
5286        // build a record batch
5287        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
5288
5289        let mut writer = ArrowWriter::try_new(vec![], batch.schema(), None).unwrap();
5290
5291        // starts empty
5292        assert_eq!(writer.in_progress_size(), 0);
5293        assert_eq!(writer.in_progress_rows(), 0);
5294        assert_eq!(writer.memory_size(), 0);
5295        assert_eq!(writer.bytes_written(), 4); // Initial header
5296        writer.write(&batch).unwrap();
5297
5298        // updated on write
5299        let initial_size = writer.in_progress_size();
5300        assert!(initial_size > 0);
5301        assert_eq!(writer.in_progress_rows(), 5);
5302        let initial_memory = writer.memory_size();
5303        assert!(initial_memory > 0);
5304        // memory estimate is larger than estimated encoded size
5305        assert!(
5306            initial_size <= initial_memory,
5307            "{initial_size} <= {initial_memory}"
5308        );
5309
5310        // updated on second write
5311        writer.write(&batch).unwrap();
5312        assert!(writer.in_progress_size() > initial_size);
5313        assert_eq!(writer.in_progress_rows(), 10);
5314        assert!(writer.memory_size() > initial_memory);
5315        assert!(
5316            writer.in_progress_size() <= writer.memory_size(),
5317            "in_progress_size {} <= memory_size {}",
5318            writer.in_progress_size(),
5319            writer.memory_size()
5320        );
5321
5322        // in progress tracking is cleared, but the overall data written is updated
5323        let pre_flush_bytes_written = writer.bytes_written();
5324        writer.flush().unwrap();
5325        assert_eq!(writer.in_progress_size(), 0);
5326        assert_eq!(writer.memory_size(), 0);
5327        assert!(writer.bytes_written() > pre_flush_bytes_written);
5328
5329        writer.close().unwrap();
5330    }
5331
5332    #[test]
5333    fn test_writer_all_null() {
5334        let a = Int32Array::from(vec![1, 2, 3, 4, 5]);
5335        let b = Int32Array::new(vec![0; 5].into(), Some(NullBuffer::new_null(5)));
5336        let batch = RecordBatch::try_from_iter(vec![
5337            ("a", Arc::new(a) as ArrayRef),
5338            ("b", Arc::new(b) as ArrayRef),
5339        ])
5340        .unwrap();
5341
5342        let mut buf = Vec::with_capacity(1024);
5343        let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), None).unwrap();
5344        writer.write(&batch).unwrap();
5345        writer.close().unwrap();
5346
5347        let bytes = Bytes::from(buf);
5348        let options = ReadOptionsBuilder::new().with_page_index().build();
5349        let reader = SerializedFileReader::new_with_options(bytes, options).unwrap();
5350        let index = reader.metadata().page_index().unwrap();
5351
5352        assert_eq!(index.num_data_pages(0, 0), Some(1)); // 1 page
5353        assert_eq!(index.num_data_pages(0, 1), Some(1)); // 1 page
5354    }
5355
5356    #[test]
5357    fn test_disabled_statistics_with_page() {
5358        let file_schema = Schema::new(vec![
5359            Field::new("a", DataType::Utf8, true),
5360            Field::new("b", DataType::Utf8, true),
5361        ]);
5362        let file_schema = Arc::new(file_schema);
5363
5364        let batch = RecordBatch::try_new(
5365            file_schema.clone(),
5366            vec![
5367                Arc::new(StringArray::from(vec!["a", "b", "c", "d"])) as _,
5368                Arc::new(StringArray::from(vec!["w", "x", "y", "z"])) as _,
5369            ],
5370        )
5371        .unwrap();
5372
5373        let props = WriterProperties::builder()
5374            .set_statistics_enabled(EnabledStatistics::None)
5375            .set_column_statistics_enabled("a".into(), EnabledStatistics::Page)
5376            .build();
5377
5378        let mut buf = Vec::with_capacity(1024);
5379        let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), Some(props)).unwrap();
5380        writer.write(&batch).unwrap();
5381
5382        let metadata = writer.close().unwrap();
5383        assert_eq!(metadata.num_row_groups(), 1);
5384        let row_group = metadata.row_group(0);
5385        assert_eq!(row_group.num_columns(), 2);
5386        // Column "a" has both offset and column index, as requested
5387        assert!(row_group.column(0).offset_index_offset().is_some());
5388        assert!(row_group.column(0).column_index_offset().is_some());
5389        // Column "b" should only have offset index
5390        assert!(row_group.column(1).offset_index_offset().is_some());
5391        assert!(row_group.column(1).column_index_offset().is_none());
5392
5393        let options = ReadOptionsBuilder::new().with_page_index().build();
5394        let reader = SerializedFileReader::new_with_options(Bytes::from(buf), options).unwrap();
5395
5396        let row_group = reader.get_row_group(0).unwrap();
5397        let a_col = row_group.metadata().column(0);
5398        let b_col = row_group.metadata().column(1);
5399
5400        // Column chunk of column "a" should have chunk level statistics
5401        if let Statistics::ByteArray(byte_array_stats) = a_col.statistics().unwrap() {
5402            let min = byte_array_stats.min_opt().unwrap();
5403            let max = byte_array_stats.max_opt().unwrap();
5404
5405            assert_eq!(min.as_bytes(), b"a");
5406            assert_eq!(max.as_bytes(), b"d");
5407        } else {
5408            panic!("expecting Statistics::ByteArray");
5409        }
5410
5411        // The column chunk for column "b" shouldn't have statistics
5412        assert!(b_col.statistics().is_none());
5413
5414        let page_index = reader.metadata().page_index().unwrap();
5415
5416        let a_idx = page_index.column_index(0, 0);
5417        assert!(
5418            matches!(a_idx, Some(ColumnIndexMetaData::BYTE_ARRAY(_))),
5419            "{a_idx:?}"
5420        );
5421        let b_idx = page_index.column_index(0, 1);
5422        assert!(b_idx.is_none(), "{b_idx:?}");
5423    }
5424
5425    #[test]
5426    fn test_disabled_statistics_with_chunk() {
5427        let file_schema = Schema::new(vec![
5428            Field::new("a", DataType::Utf8, true),
5429            Field::new("b", DataType::Utf8, true),
5430        ]);
5431        let file_schema = Arc::new(file_schema);
5432
5433        let batch = RecordBatch::try_new(
5434            file_schema.clone(),
5435            vec![
5436                Arc::new(StringArray::from(vec!["a", "b", "c", "d"])) as _,
5437                Arc::new(StringArray::from(vec!["w", "x", "y", "z"])) as _,
5438            ],
5439        )
5440        .unwrap();
5441
5442        let props = WriterProperties::builder()
5443            .set_statistics_enabled(EnabledStatistics::None)
5444            .set_column_statistics_enabled("a".into(), EnabledStatistics::Chunk)
5445            .build();
5446
5447        let mut buf = Vec::with_capacity(1024);
5448        let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), Some(props)).unwrap();
5449        writer.write(&batch).unwrap();
5450
5451        let metadata = writer.close().unwrap();
5452        assert_eq!(metadata.num_row_groups(), 1);
5453        let row_group = metadata.row_group(0);
5454        assert_eq!(row_group.num_columns(), 2);
5455        // Column "a" should only have offset index
5456        assert!(row_group.column(0).offset_index_offset().is_some());
5457        assert!(row_group.column(0).column_index_offset().is_none());
5458        // Column "b" should only have offset index
5459        assert!(row_group.column(1).offset_index_offset().is_some());
5460        assert!(row_group.column(1).column_index_offset().is_none());
5461
5462        let options = ReadOptionsBuilder::new().with_page_index().build();
5463        let reader = SerializedFileReader::new_with_options(Bytes::from(buf), options).unwrap();
5464
5465        let row_group = reader.get_row_group(0).unwrap();
5466        let a_col = row_group.metadata().column(0);
5467        let b_col = row_group.metadata().column(1);
5468
5469        // Column chunk of column "a" should have chunk level statistics
5470        if let Statistics::ByteArray(byte_array_stats) = a_col.statistics().unwrap() {
5471            let min = byte_array_stats.min_opt().unwrap();
5472            let max = byte_array_stats.max_opt().unwrap();
5473
5474            assert_eq!(min.as_bytes(), b"a");
5475            assert_eq!(max.as_bytes(), b"d");
5476        } else {
5477            panic!("expecting Statistics::ByteArray");
5478        }
5479
5480        // The column chunk for column "b"  shouldn't have statistics
5481        assert!(b_col.statistics().is_none());
5482
5483        let page_index = reader.metadata().page_index().unwrap();
5484
5485        let a_idx = page_index.column_index(0, 0);
5486        assert!(a_idx.is_none(), "{a_idx:?}");
5487        let b_idx = page_index.column_index(0, 1);
5488        assert!(b_idx.is_none(), "{b_idx:?}");
5489    }
5490
5491    #[test]
5492    fn test_arrow_writer_skip_metadata() {
5493        let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
5494        let file_schema = Arc::new(batch_schema.clone());
5495
5496        let batch = RecordBatch::try_new(
5497            Arc::new(batch_schema),
5498            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5499        )
5500        .unwrap();
5501        let skip_options = ArrowWriterOptions::new().with_skip_arrow_metadata(true);
5502
5503        let mut buf = Vec::with_capacity(1024);
5504        let mut writer =
5505            ArrowWriter::try_new_with_options(&mut buf, file_schema.clone(), skip_options).unwrap();
5506        writer.write(&batch).unwrap();
5507        writer.close().unwrap();
5508
5509        let bytes = Bytes::from(buf);
5510        let reader_builder = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap();
5511        assert_eq!(file_schema, *reader_builder.schema());
5512        if let Some(key_value_metadata) = reader_builder
5513            .metadata()
5514            .file_metadata()
5515            .key_value_metadata()
5516        {
5517            assert!(
5518                !key_value_metadata
5519                    .iter()
5520                    .any(|kv| kv.key.as_str() == ARROW_SCHEMA_META_KEY)
5521            );
5522        }
5523    }
5524
5525    #[test]
5526    fn test_arrow_writer_skip_path_in_schema() {
5527        let batch_schema = Schema::new(vec![Field::new("int32", DataType::Int32, false)]);
5528        let file_schema = Arc::new(batch_schema.clone());
5529
5530        let batch = RecordBatch::try_new(
5531            Arc::new(batch_schema),
5532            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5533        )
5534        .unwrap();
5535
5536        // default options should still write path_in_schema
5537        let skip_options = ArrowWriterOptions::new();
5538
5539        let mut buf = Vec::with_capacity(1024);
5540        let mut writer =
5541            ArrowWriter::try_new_with_options(&mut buf, file_schema.clone(), skip_options).unwrap();
5542        writer.write(&batch).unwrap();
5543        writer.close().unwrap();
5544
5545        // override to not write path_in_schema
5546        let skip_options = ArrowWriterOptions::new().with_properties(
5547            WriterProperties::builder()
5548                .set_write_path_in_schema(false)
5549                .build(),
5550        );
5551
5552        let mut buf2 = Vec::with_capacity(1024);
5553        let mut writer =
5554            ArrowWriter::try_new_with_options(&mut buf2, file_schema.clone(), skip_options)
5555                .unwrap();
5556        writer.write(&batch).unwrap();
5557        writer.close().unwrap();
5558
5559        // buf2 should be a bit smaller due to lack of path_in_schema
5560        assert!(buf.len() > buf2.len());
5561    }
5562
5563    #[test]
5564    fn mismatched_schemas() {
5565        let batch_schema = Schema::new(vec![Field::new("count", DataType::Int32, false)]);
5566        let file_schema = Arc::new(Schema::new(vec![Field::new(
5567            "temperature",
5568            DataType::Float64,
5569            false,
5570        )]));
5571
5572        let batch = RecordBatch::try_new(
5573            Arc::new(batch_schema),
5574            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5575        )
5576        .unwrap();
5577
5578        let mut buf = Vec::with_capacity(1024);
5579        let mut writer = ArrowWriter::try_new(&mut buf, file_schema.clone(), None).unwrap();
5580
5581        let err = writer.write(&batch).unwrap_err().to_string();
5582        assert_eq!(
5583            err,
5584            "Arrow: Incompatible type. Field 'temperature' has type Float64, array has type Int32"
5585        );
5586    }
5587
5588    #[test]
5589    // https://github.com/apache/arrow-rs/issues/6988
5590    fn test_roundtrip_empty_schema() {
5591        // create empty record batch with empty schema
5592        let empty_batch = RecordBatch::try_new_with_options(
5593            Arc::new(Schema::empty()),
5594            vec![],
5595            &RecordBatchOptions::default().with_row_count(Some(0)),
5596        )
5597        .unwrap();
5598
5599        // write to parquet
5600        let mut parquet_bytes: Vec<u8> = Vec::new();
5601        let mut writer =
5602            ArrowWriter::try_new(&mut parquet_bytes, empty_batch.schema(), None).unwrap();
5603        writer.write(&empty_batch).unwrap();
5604        writer.close().unwrap();
5605
5606        // read from parquet
5607        let bytes = Bytes::from(parquet_bytes);
5608        let reader = ParquetRecordBatchReaderBuilder::try_new(bytes).unwrap();
5609        assert_eq!(reader.schema(), &empty_batch.schema());
5610        let batches: Vec<_> = reader
5611            .build()
5612            .unwrap()
5613            .collect::<ArrowResult<Vec<_>>>()
5614            .unwrap();
5615        assert_eq!(batches.len(), 0);
5616    }
5617
5618    #[test]
5619    fn test_page_stats_not_written_by_default() {
5620        let string_field = Field::new("a", DataType::Utf8, false);
5621        let schema = Schema::new(vec![string_field]);
5622        let raw_string_values = vec!["Blart Versenwald III"];
5623        let string_values = StringArray::from(raw_string_values.clone());
5624        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(string_values)]).unwrap();
5625
5626        let props = WriterProperties::builder()
5627            .set_statistics_enabled(EnabledStatistics::Page)
5628            .set_dictionary_enabled(false)
5629            .set_encoding(Encoding::PLAIN)
5630            .set_compression(crate::basic::Compression::UNCOMPRESSED)
5631            .build();
5632
5633        let file = roundtrip_opts(&batch, props);
5634
5635        // read file and decode page headers
5636        // Note: use the thrift API as there is no Rust API to access the statistics in the page headers
5637
5638        // decode first page header
5639        let first_page = &file[4..];
5640        let mut prot = ThriftSliceInputProtocol::new(first_page);
5641        let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5642        let stats = hdr.data_page_header.unwrap().statistics;
5643
5644        assert!(stats.is_none());
5645    }
5646
5647    #[test]
5648    fn test_page_stats_when_enabled() {
5649        let string_field = Field::new("a", DataType::Utf8, false);
5650        let schema = Schema::new(vec![string_field]);
5651        let raw_string_values = vec!["Blart Versenwald III", "Andrew Lamb"];
5652        let string_values = StringArray::from(raw_string_values.clone());
5653        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(string_values)]).unwrap();
5654
5655        let props = WriterProperties::builder()
5656            .set_statistics_enabled(EnabledStatistics::Page)
5657            .set_dictionary_enabled(false)
5658            .set_encoding(Encoding::PLAIN)
5659            .set_write_page_header_statistics(true)
5660            .set_compression(crate::basic::Compression::UNCOMPRESSED)
5661            .build();
5662
5663        let file = roundtrip_opts(&batch, props);
5664
5665        // read file and decode page headers
5666        // Note: use the thrift API as there is no Rust API to access the statistics in the page headers
5667
5668        // decode first page header
5669        let first_page = &file[4..];
5670        let mut prot = ThriftSliceInputProtocol::new(first_page);
5671        let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5672        let stats = hdr.data_page_header.unwrap().statistics;
5673
5674        let stats = stats.unwrap();
5675        // check that min/max were actually written to the page
5676        assert!(stats.is_max_value_exact.unwrap());
5677        assert!(stats.is_min_value_exact.unwrap());
5678        assert_eq!(stats.max_value.unwrap(), b"Blart Versenwald III");
5679        assert_eq!(stats.min_value.unwrap(), b"Andrew Lamb");
5680    }
5681
5682    #[test]
5683    fn test_page_stats_truncation() {
5684        let string_field = Field::new("a", DataType::Utf8, false);
5685        let binary_field = Field::new("b", DataType::Binary, false);
5686        let schema = Schema::new(vec![string_field, binary_field]);
5687
5688        let raw_string_values = vec!["Blart Versenwald III"];
5689        let raw_binary_values = [b"Blart Versenwald III".to_vec()];
5690        let raw_binary_value_refs = raw_binary_values
5691            .iter()
5692            .map(|x| x.as_slice())
5693            .collect::<Vec<_>>();
5694
5695        let string_values = StringArray::from(raw_string_values.clone());
5696        let binary_values = BinaryArray::from(raw_binary_value_refs);
5697        let batch = RecordBatch::try_new(
5698            Arc::new(schema),
5699            vec![Arc::new(string_values), Arc::new(binary_values)],
5700        )
5701        .unwrap();
5702
5703        let props = WriterProperties::builder()
5704            .set_statistics_truncate_length(Some(2))
5705            .set_dictionary_enabled(false)
5706            .set_encoding(Encoding::PLAIN)
5707            .set_write_page_header_statistics(true)
5708            .set_compression(crate::basic::Compression::UNCOMPRESSED)
5709            .build();
5710
5711        let file = roundtrip_opts(&batch, props);
5712
5713        // read file and decode page headers
5714        // Note: use the thrift API as there is no Rust API to access the statistics in the page headers
5715
5716        // decode first page header
5717        let first_page = &file[4..];
5718        let mut prot = ThriftSliceInputProtocol::new(first_page);
5719        let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5720        let stats = hdr.data_page_header.unwrap().statistics;
5721        assert!(stats.is_some());
5722        let stats = stats.unwrap();
5723        // check that min/max were properly truncated
5724        assert!(!stats.is_max_value_exact.unwrap());
5725        assert!(!stats.is_min_value_exact.unwrap());
5726        assert_eq!(stats.max_value.unwrap(), b"Bm");
5727        assert_eq!(stats.min_value.unwrap(), b"Bl");
5728
5729        // check second page now
5730        let second_page = &prot.as_slice()[hdr.compressed_page_size as usize..];
5731        let mut prot = ThriftSliceInputProtocol::new(second_page);
5732        let hdr = PageHeader::read_thrift(&mut prot).unwrap();
5733        let stats = hdr.data_page_header.unwrap().statistics;
5734        assert!(stats.is_some());
5735        let stats = stats.unwrap();
5736        // check that min/max were properly truncated
5737        assert!(!stats.is_max_value_exact.unwrap());
5738        assert!(!stats.is_min_value_exact.unwrap());
5739        assert_eq!(stats.max_value.unwrap(), b"Bm");
5740        assert_eq!(stats.min_value.unwrap(), b"Bl");
5741    }
5742
5743    #[test]
5744    fn test_page_encoding_statistics_roundtrip() {
5745        let batch_schema = Schema::new(vec![Field::new(
5746            "int32",
5747            arrow_schema::DataType::Int32,
5748            false,
5749        )]);
5750
5751        let batch = RecordBatch::try_new(
5752            Arc::new(batch_schema.clone()),
5753            vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as _],
5754        )
5755        .unwrap();
5756
5757        let mut file: File = tempfile::tempfile().unwrap();
5758        let mut writer = ArrowWriter::try_new(&mut file, Arc::new(batch_schema), None).unwrap();
5759        writer.write(&batch).unwrap();
5760        let file_metadata = writer.close().unwrap();
5761
5762        assert_eq!(file_metadata.num_row_groups(), 1);
5763        assert_eq!(file_metadata.row_group(0).num_columns(), 1);
5764        assert!(
5765            file_metadata
5766                .row_group(0)
5767                .column(0)
5768                .page_encoding_stats()
5769                .is_some()
5770        );
5771        let chunk_page_stats = file_metadata
5772            .row_group(0)
5773            .column(0)
5774            .page_encoding_stats()
5775            .unwrap();
5776
5777        // check that the read metadata is also correct
5778        let options = ReadOptionsBuilder::new()
5779            .with_page_index()
5780            .with_encoding_stats_as_mask(false)
5781            .build();
5782        let reader = SerializedFileReader::new_with_options(file, options).unwrap();
5783
5784        let rowgroup = reader.get_row_group(0).expect("row group missing");
5785        assert_eq!(rowgroup.num_columns(), 1);
5786        let column = rowgroup.metadata().column(0);
5787        assert!(column.page_encoding_stats().is_some());
5788        let file_page_stats = column.page_encoding_stats().unwrap();
5789        assert_eq!(chunk_page_stats, file_page_stats);
5790    }
5791
5792    #[test]
5793    #[cfg_attr(miri, ignore)] // Takes too long
5794    fn test_different_dict_page_size_limit() {
5795        let array = Arc::new(Int64Array::from_iter(0..1024 * 1024));
5796        let schema = Arc::new(Schema::new(vec![
5797            Field::new("col0", arrow_schema::DataType::Int64, false),
5798            Field::new("col1", arrow_schema::DataType::Int64, false),
5799        ]));
5800        let batch =
5801            arrow_array::RecordBatch::try_new(schema.clone(), vec![array.clone(), array]).unwrap();
5802
5803        let props = WriterProperties::builder()
5804            .set_dictionary_page_size_limit(1024 * 1024)
5805            .set_column_dictionary_page_size_limit(ColumnPath::from("col1"), 1024 * 1024 * 4)
5806            .build();
5807        let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5808        writer.write(&batch).unwrap();
5809        let data = Bytes::from(writer.into_inner().unwrap());
5810
5811        let mut metadata = ParquetMetaDataReader::new();
5812        metadata.try_parse(&data).unwrap();
5813        let metadata = metadata.finish().unwrap();
5814        let col0_meta = metadata.row_group(0).column(0);
5815        let col1_meta = metadata.row_group(0).column(1);
5816
5817        let get_dict_page_size = move |meta: &ColumnChunkMetaData| {
5818            let mut reader =
5819                SerializedPageReader::new(Arc::new(data.clone()), meta, 0, None).unwrap();
5820            let page = reader.get_next_page().unwrap().unwrap();
5821            match page {
5822                Page::DictionaryPage { buf, .. } => buf.len(),
5823                _ => panic!("expected DictionaryPage"),
5824            }
5825        };
5826
5827        assert_eq!(get_dict_page_size(col0_meta), 1024 * 1024);
5828        assert_eq!(get_dict_page_size(col1_meta), 1024 * 1024 * 4);
5829    }
5830
5831    #[test]
5832    #[cfg_attr(miri, ignore)] // Takes too long
5833    fn test_arrow_writer_granular_mode_roundtrip() {
5834        // Granular mode subdivides chunks and writes more pages than the
5835        // default batched path. Make sure the data we write back is
5836        // bit-identical to what went in — page-count assertions elsewhere
5837        // only prove pages were cut, not that the encoded data is correct.
5838        //
5839        // Mix value sizes so that the cumulative-byte-budget cutoff
5840        // lands mid-chunk, exercising both batched and granular paths
5841        // within the same `write_batch_internal` call.
5842        let small = "tiny".to_string();
5843        let big = "x".repeat(64 * 1024);
5844        let strings: Vec<String> = (0..256)
5845            .map(|i| {
5846                if i % 16 == 0 {
5847                    big.clone()
5848                } else {
5849                    small.clone()
5850                }
5851            })
5852            .collect();
5853
5854        let schema = Arc::new(Schema::new(vec![Field::new(
5855            "col",
5856            ArrowDataType::Utf8,
5857            false,
5858        )]));
5859        let batch = RecordBatch::try_new(
5860            schema.clone(),
5861            vec![Arc::new(StringArray::from(strings.clone())) as _],
5862        )
5863        .unwrap();
5864
5865        let props = WriterProperties::builder()
5866            .set_dictionary_enabled(false)
5867            .set_data_page_size_limit(16 * 1024)
5868            .build();
5869        let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5870        writer.write(&batch).unwrap();
5871        let data = Bytes::from(writer.into_inner().unwrap());
5872
5873        let mut reader = ParquetRecordBatchReader::try_new(data, 1024).unwrap();
5874        let read = reader.next().unwrap().unwrap();
5875        assert!(reader.next().is_none(), "expected one batch");
5876        let col = read
5877            .column(0)
5878            .as_any()
5879            .downcast_ref::<StringArray>()
5880            .unwrap();
5881        assert_eq!(col.len(), strings.len());
5882        for (i, expected) in strings.iter().enumerate() {
5883            assert_eq!(
5884                col.value(i),
5885                expected.as_str(),
5886                "value mismatch at index {i}"
5887            );
5888        }
5889    }
5890
5891    #[test]
5892    fn test_arrow_writer_all_null_string_column() {
5893        // The `LevelDataRef::value_count` Uniform branch with
5894        // `value != max_def` (entirely-null chunk) must return 0 so the
5895        // sub-batch sizer short-circuits to batch mode without trying
5896        // to estimate byte budgets for non-existent values.
5897        let num_rows = 1024;
5898        let schema = Arc::new(Schema::new(vec![Field::new(
5899            "col",
5900            ArrowDataType::Utf8,
5901            true,
5902        )]));
5903        let nulls: Vec<Option<&str>> = vec![None; num_rows];
5904        let batch = RecordBatch::try_new(
5905            schema.clone(),
5906            vec![Arc::new(StringArray::from(nulls)) as _],
5907        )
5908        .unwrap();
5909
5910        let props = WriterProperties::builder()
5911            .set_dictionary_enabled(false)
5912            .set_data_page_size_limit(16 * 1024)
5913            .build();
5914        let mut writer = ArrowWriter::try_new(Vec::new(), schema, Some(props)).unwrap();
5915        writer.write(&batch).unwrap();
5916        let data = Bytes::from(writer.into_inner().unwrap());
5917
5918        // Re-parse the file: row group has one column, every row is
5919        // null, all data pages report `num_rows / page_count` rows.
5920        let mut metadata = ParquetMetaDataReader::new();
5921        metadata.try_parse(&data).unwrap();
5922        let metadata = metadata.finish().unwrap();
5923        let row_group = metadata.row_group(0);
5924        let col_meta = row_group.column(0);
5925        assert_eq!(row_group.num_rows() as usize, num_rows);
5926        // Statistics record `null_count = num_rows` — proves every value
5927        // was written as null.
5928        if let Some(stats) = col_meta.statistics() {
5929            assert_eq!(
5930                stats.null_count_opt().unwrap_or(0) as usize,
5931                num_rows,
5932                "expected all-null column to report null_count = num_rows"
5933            );
5934        }
5935
5936        let mut reader =
5937            SerializedPageReader::new(Arc::new(data.clone()), col_meta, num_rows, None).unwrap();
5938        let mut total_values = 0u32;
5939        while let Some(page) = reader.get_next_page().unwrap() {
5940            if matches!(page, Page::DataPage { .. } | Page::DataPageV2 { .. }) {
5941                total_values += page.num_values();
5942            }
5943        }
5944        assert_eq!(
5945            total_values as usize, num_rows,
5946            "expected every level position to be represented in some page"
5947        );
5948    }
5949
5950    struct WriteBatchesShape {
5951        num_batches: usize,
5952        rows_per_batch: usize,
5953        row_size: usize,
5954    }
5955
5956    /// Helper function to write batches with the provided `WriteBatchesShape` into an `ArrowWriter`
5957    fn write_batches(
5958        WriteBatchesShape {
5959            num_batches,
5960            rows_per_batch,
5961            row_size,
5962        }: WriteBatchesShape,
5963        props: WriterProperties,
5964    ) -> ParquetRecordBatchReaderBuilder<File> {
5965        let schema = Arc::new(Schema::new(vec![Field::new(
5966            "str",
5967            ArrowDataType::Utf8,
5968            false,
5969        )]));
5970        let file = tempfile::tempfile().unwrap();
5971        let mut writer =
5972            ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
5973
5974        for batch_idx in 0..num_batches {
5975            let strings: Vec<String> = (0..rows_per_batch)
5976                .map(|i| format!("{:0>width$}", batch_idx * 10 + i, width = row_size))
5977                .collect();
5978            let array = StringArray::from(strings);
5979            let batch = RecordBatch::try_new(schema.clone(), vec![Arc::new(array)]).unwrap();
5980            writer.write(&batch).unwrap();
5981        }
5982        writer.close().unwrap();
5983        ParquetRecordBatchReaderBuilder::try_new(file).unwrap()
5984    }
5985
5986    #[test]
5987    // When both limits are None, all data should go into a single row group
5988    fn test_row_group_limit_none_writes_single_row_group() {
5989        let props = WriterProperties::builder()
5990            .set_max_row_group_row_count(None)
5991            .set_max_row_group_bytes(None)
5992            .build();
5993
5994        let builder = write_batches(
5995            WriteBatchesShape {
5996                num_batches: 1,
5997                rows_per_batch: 1000,
5998                row_size: 4,
5999            },
6000            props,
6001        );
6002
6003        assert_eq!(
6004            &row_group_sizes(builder.metadata()),
6005            &[1000],
6006            "With no limits, all rows should be in a single row group"
6007        );
6008    }
6009
6010    #[test]
6011    // When only max_row_group_size is set, respect the row limit
6012    fn test_row_group_limit_rows_only() {
6013        let props = WriterProperties::builder()
6014            .set_max_row_group_row_count(Some(300))
6015            .set_max_row_group_bytes(None)
6016            .build();
6017
6018        let builder = write_batches(
6019            WriteBatchesShape {
6020                num_batches: 1,
6021                rows_per_batch: 1000,
6022                row_size: 4,
6023            },
6024            props,
6025        );
6026
6027        assert_eq!(
6028            &row_group_sizes(builder.metadata()),
6029            &[300, 300, 300, 100],
6030            "Row groups should be split by row count"
6031        );
6032    }
6033
6034    #[test]
6035    #[cfg_attr(miri, ignore)] // Takes too long
6036    // A row limit far smaller than the batch splits it many times over; the split must not
6037    // consume stack proportional to the number of row groups.
6038    fn test_row_group_limit_rows_only_many_splits() {
6039        let props = WriterProperties::builder()
6040            .set_max_row_group_row_count(Some(1))
6041            .set_max_row_group_bytes(None)
6042            .build();
6043
6044        let rows = 50_000;
6045        let builder = write_batches(
6046            WriteBatchesShape {
6047                num_batches: 1,
6048                rows_per_batch: rows,
6049                row_size: 4,
6050            },
6051            props,
6052        );
6053
6054        let sizes = row_group_sizes(builder.metadata());
6055        assert_eq!(sizes.len(), rows, "Every row should get its own row group");
6056        assert_eq!(
6057            sizes.iter().sum::<i64>(),
6058            rows as i64,
6059            "Total rows should be preserved"
6060        );
6061    }
6062
6063    #[test]
6064    // When only max_row_group_bytes is set, respect the byte limit
6065    fn test_row_group_limit_bytes_only() {
6066        let props = WriterProperties::builder()
6067            .set_max_row_group_row_count(None)
6068            // Set byte limit to approximately fit ~30 rows worth of data (~100 bytes each)
6069            .set_max_row_group_bytes(Some(3500))
6070            .build();
6071
6072        let builder = write_batches(
6073            WriteBatchesShape {
6074                num_batches: 10,
6075                rows_per_batch: 10,
6076                row_size: 100,
6077            },
6078            props,
6079        );
6080
6081        let sizes = row_group_sizes(builder.metadata());
6082
6083        assert!(
6084            sizes.len() > 1,
6085            "Should have multiple row groups due to byte limit, got {sizes:?}",
6086        );
6087
6088        let total_rows: i64 = sizes.iter().sum();
6089        assert_eq!(total_rows, 100, "Total rows should be preserved");
6090    }
6091
6092    #[test]
6093    // If an in-progress row group is already oversized, it should be flushed before writing more.
6094    fn test_row_group_limit_bytes_flushes_when_current_group_already_too_large() {
6095        let schema = Arc::new(Schema::new(vec![Field::new(
6096            "str",
6097            ArrowDataType::Utf8,
6098            false,
6099        )]));
6100        let file = tempfile::tempfile().unwrap();
6101
6102        // Start with no byte limit so we can intentionally build an oversized in-progress row group.
6103        let props = WriterProperties::builder()
6104            .set_max_row_group_row_count(None)
6105            .set_max_row_group_bytes(None)
6106            .build();
6107        let mut writer =
6108            ArrowWriter::try_new(file.try_clone().unwrap(), schema.clone(), Some(props)).unwrap();
6109
6110        let first_array = StringArray::from(
6111            (0..10)
6112                .map(|i| format!("{i:0>100}"))
6113                .collect::<Vec<String>>(),
6114        );
6115        let first_batch =
6116            RecordBatch::try_new(schema.clone(), vec![Arc::new(first_array)]).unwrap();
6117        writer.write(&first_batch).unwrap();
6118        assert_eq!(writer.in_progress_rows(), 10);
6119
6120        // Tighten the limit below the current in-progress bytes to exercise:
6121        // `if current_bytes >= max_bytes { self.flush()?; ... }`
6122        writer.max_row_group_bytes = Some(1);
6123
6124        let second_array = StringArray::from(vec!["x".to_string()]);
6125        let second_batch =
6126            RecordBatch::try_new(schema.clone(), vec![Arc::new(second_array)]).unwrap();
6127        writer.write(&second_batch).unwrap();
6128        writer.close().unwrap();
6129        let builder = ParquetRecordBatchReaderBuilder::try_new(file).unwrap();
6130
6131        assert_eq!(
6132            &row_group_sizes(builder.metadata()),
6133            &[10, 1],
6134            "The second write should flush an oversized in-progress row group first",
6135        );
6136    }
6137
6138    #[test]
6139    // When both limits are set, the row limit triggers first
6140    fn test_row_group_limit_both_row_wins_single_batch() {
6141        let props = WriterProperties::builder()
6142            .set_max_row_group_row_count(Some(200)) // Will trigger at 200 rows
6143            .set_max_row_group_bytes(Some(1024 * 1024)) // 1MB - won't trigger for small int data
6144            .build();
6145
6146        let builder = write_batches(
6147            WriteBatchesShape {
6148                num_batches: 1,
6149                row_size: 4,
6150                rows_per_batch: 1000,
6151            },
6152            props,
6153        );
6154
6155        assert_eq!(
6156            &row_group_sizes(builder.metadata()),
6157            &[200, 200, 200, 200, 200],
6158            "Row limit should trigger before byte limit"
6159        );
6160    }
6161
6162    #[test]
6163    // When both limits are set, the row limit triggers first
6164    fn test_row_group_limit_both_row_wins_multiple_batches() {
6165        let props = WriterProperties::builder()
6166            .set_max_row_group_row_count(Some(5)) // Will trigger every 5 rows
6167            .set_max_row_group_bytes(Some(9999)) // Won't trigger
6168            .build();
6169
6170        let builder = write_batches(
6171            WriteBatchesShape {
6172                num_batches: 10,
6173                rows_per_batch: 10,
6174                row_size: 100,
6175            },
6176            props,
6177        );
6178
6179        assert_eq!(
6180            &row_group_sizes(builder.metadata()),
6181            &[5; 20],
6182            "Row limit should trigger before byte limit"
6183        );
6184    }
6185
6186    #[test]
6187    // When both limits are set, the byte limit triggers first
6188    fn test_row_group_limit_both_bytes_wins() {
6189        let props = WriterProperties::builder()
6190            .set_max_row_group_row_count(Some(1000)) // Won't trigger for 100 rows
6191            .set_max_row_group_bytes(Some(3500)) // Will trigger at ~30-35 rows
6192            .build();
6193
6194        let builder = write_batches(
6195            WriteBatchesShape {
6196                num_batches: 10,
6197                rows_per_batch: 10,
6198                row_size: 100,
6199            },
6200            props,
6201        );
6202
6203        let sizes = row_group_sizes(builder.metadata());
6204
6205        assert!(
6206            sizes.len() > 1,
6207            "Byte limit should trigger before row limit, got {sizes:?}",
6208        );
6209
6210        assert!(
6211            sizes.iter().all(|&s| s < 1000),
6212            "No row group should hit the row limit"
6213        );
6214
6215        let total_rows: i64 = sizes.iter().sum();
6216        assert_eq!(total_rows, 100, "Total rows should be preserved");
6217    }
6218
6219    #[test]
6220    // Both limits can apply to the same batch: the row limit trims it to 5 rows, and the
6221    // byte limit then trims those 5 down to 4.
6222    fn test_row_group_limit_both_apply_to_same_batch() {
6223        let props = WriterProperties::builder()
6224            .set_max_row_group_row_count(Some(15))
6225            .set_max_row_group_bytes(Some(1500))
6226            .build();
6227
6228        let builder = write_batches(
6229            WriteBatchesShape {
6230                num_batches: 2,
6231                rows_per_batch: 10,
6232                row_size: 100,
6233            },
6234            props,
6235        );
6236
6237        assert_eq!(
6238            &row_group_sizes(builder.metadata()),
6239            &[14, 6],
6240            "Byte limit should still apply to a batch the row limit already split"
6241        );
6242    }
6243
6244    #[test]
6245    fn arrow_column_chunk_close_mut_drops_column_index() {
6246        use crate::arrow::ArrowSchemaConverter;
6247        use crate::file::writer::SerializedFileWriter;
6248
6249        let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)]));
6250        let props = Arc::new(
6251            WriterProperties::builder()
6252                .set_statistics_enabled(EnabledStatistics::Page)
6253                .build(),
6254        );
6255        let parquet_schema = ArrowSchemaConverter::new()
6256            .with_coerce_types(props.coerce_types())
6257            .convert(&schema)
6258            .unwrap();
6259
6260        let mut buf = Vec::with_capacity(1024);
6261        let mut writer =
6262            SerializedFileWriter::new(&mut buf, parquet_schema.root_schema_ptr(), props.clone())
6263                .unwrap();
6264
6265        let factory = ArrowRowGroupWriterFactory::new(&writer, Arc::clone(&schema));
6266        let mut col_writers = factory.create_column_writers(0).unwrap();
6267        let arr: ArrayRef = Arc::new(Int32Array::from_iter_values(0..64));
6268        for leaves in compute_leaves(schema.field(0), &arr).unwrap() {
6269            col_writers[0].write(&leaves).unwrap();
6270        }
6271        let mut chunk = col_writers.pop().unwrap().close().unwrap();
6272
6273        // Immutable accessor exposes the close result produced at close time.
6274        assert!(
6275            chunk.close().column_index.is_some(),
6276            "EnabledStatistics::Page should produce a column_index"
6277        );
6278
6279        // Mutable accessor lets callers drop the page-level index before append.
6280        chunk.close_mut().column_index = None;
6281        assert!(chunk.close().column_index.is_none());
6282
6283        let mut rg = writer.next_row_group().unwrap();
6284        chunk.append_to_row_group(&mut rg).unwrap();
6285        rg.close().unwrap();
6286        let file_meta = writer.close().unwrap();
6287
6288        // After dropping column_index, the resulting file records no column
6289        // index offset/length for this chunk.
6290        let cc = file_meta.row_group(0).column(0);
6291        assert!(cc.column_index_range().is_none());
6292    }
6293
6294    /// Writes a single-column RecordBatch to an in-memory Parquet buffer.
6295    fn write_column_to_bytes(array: ArrayRef) -> Bytes {
6296        let schema = Arc::new(Schema::new(vec![Field::new(
6297            "col",
6298            array.data_type().clone(),
6299            true,
6300        )]));
6301        let buf = get_bytes_after_close(
6302            schema.clone(),
6303            &RecordBatch::try_new(schema, vec![array]).unwrap(),
6304        );
6305        Bytes::from(buf)
6306    }
6307
6308    /// Reads column 0 from a single-row-group Parquet buffer, projecting it with the given schema.
6309    /// Passing a flat schema when the buffer was written from a REE array lets callers decode
6310    /// the physical values without the run-end encoding wrapper.
6311    fn read_column_with_schema(bytes: Bytes, schema: SchemaRef) -> ArrayRef {
6312        let opts = crate::arrow::arrow_reader::ArrowReaderOptions::new().with_schema(schema);
6313        ParquetRecordBatchReaderBuilder::try_new_with_options(bytes, opts)
6314            .unwrap()
6315            .build()
6316            .unwrap()
6317            .next()
6318            .unwrap()
6319            .unwrap()
6320            .column(0)
6321            .clone()
6322    }
6323
6324    fn ree_write_read_roundtrip(ree: ArrayRef, flat: ArrayRef) {
6325        let flat_schema = Arc::new(Schema::new(vec![Field::new(
6326            "col",
6327            flat.data_type().clone(),
6328            true,
6329        )]));
6330        let ree_bytes = write_column_to_bytes(ree);
6331        let flat_bytes = write_column_to_bytes(flat.clone());
6332        assert_eq!(
6333            ree_bytes, flat_bytes,
6334            "REE and flat bytes should be identical"
6335        );
6336
6337        let decoded_ree = read_column_with_schema(ree_bytes, flat_schema.clone());
6338        let decoded_flat = read_column_with_schema(flat_bytes, flat_schema);
6339
6340        assert_eq!(decoded_ree.as_ref(), flat.as_ref());
6341        assert_eq!(decoded_ree.as_ref(), decoded_flat.as_ref());
6342    }
6343
6344    #[test]
6345    fn ree_string() {
6346        let ree: ArrayRef = Arc::new(
6347            [Some("a"), Some("a"), None, Some("b"), Some("b")]
6348                .into_iter()
6349                .collect::<Int32RunArray>(),
6350        );
6351        let flat: ArrayRef = Arc::new(StringArray::from(vec![
6352            Some("a"),
6353            Some("a"),
6354            None,
6355            Some("b"),
6356            Some("b"),
6357        ]));
6358        ree_write_read_roundtrip(ree, flat);
6359    }
6360
6361    #[test]
6362    fn ree_int32() {
6363        let mut b = PrimitiveRunBuilder::<Int32Type, Int32Type>::new();
6364        for v in [Some(1), Some(1), None, Some(2), Some(2)] {
6365            b.append_option(v);
6366        }
6367        let ree: ArrayRef = Arc::new(b.finish());
6368        let flat: ArrayRef = Arc::new(Int32Array::from(vec![
6369            Some(1),
6370            Some(1),
6371            None,
6372            Some(2),
6373            Some(2),
6374        ]));
6375        ree_write_read_roundtrip(ree, flat);
6376    }
6377
6378    #[test]
6379    fn ree_bool() {
6380        // run_ends [3, 5, 7] → [T,T,T, null,null, F,F]
6381        let ree: ArrayRef = Arc::new(
6382            RunArray::try_new(
6383                &Int32Array::from(vec![3, 5, 7]),
6384                &BooleanArray::from(vec![Some(true), None, Some(false)]),
6385            )
6386            .unwrap(),
6387        );
6388        let flat: ArrayRef = Arc::new(BooleanArray::from(vec![
6389            Some(true),
6390            Some(true),
6391            Some(true),
6392            None,
6393            None,
6394            Some(false),
6395            Some(false),
6396        ]));
6397        ree_write_read_roundtrip(ree, flat);
6398    }
6399
6400    #[test]
6401    fn ree_fixed_size_binary() {
6402        let mk = |vals: &[Option<&[u8]>]| -> FixedSizeBinaryArray {
6403            let mut b = FixedSizeBinaryBuilder::new(2);
6404            for v in vals {
6405                match v {
6406                    Some(x) => b.append_value(x).unwrap(),
6407                    None => b.append_null(),
6408                }
6409            }
6410            b.finish()
6411        };
6412        // run_ends [2, 4, 6] → [aa,aa, null,null, bb,bb]
6413        let ree: ArrayRef = Arc::new(
6414            RunArray::try_new(
6415                &Int32Array::from(vec![2, 4, 6]),
6416                &mk(&[Some(b"aa"), None, Some(b"bb")]),
6417            )
6418            .unwrap(),
6419        );
6420        let flat: ArrayRef = Arc::new(mk(&[
6421            Some(b"aa"),
6422            Some(b"aa"),
6423            None,
6424            None,
6425            Some(b"bb"),
6426            Some(b"bb"),
6427        ]));
6428        ree_write_read_roundtrip(ree, flat);
6429    }
6430
6431    #[test]
6432    fn ree_single_run() {
6433        let ree: ArrayRef = Arc::new(["x", "x", "x"].into_iter().collect::<Int32RunArray>());
6434        let flat: ArrayRef = Arc::new(StringArray::from(vec!["x", "x", "x"]));
6435        ree_write_read_roundtrip(ree, flat);
6436    }
6437
6438    #[test]
6439    fn ree_float32() {
6440        // run_ends [2, 4, 5] → [1.0, 1.0, null, null, 2.5]
6441        let ree: ArrayRef = Arc::new(
6442            RunArray::try_new(
6443                &Int32Array::from(vec![2, 4, 5]),
6444                &Float32Array::from(vec![Some(1.0_f32), None, Some(2.5_f32)]),
6445            )
6446            .unwrap(),
6447        );
6448        let flat: ArrayRef = Arc::new(Float32Array::from(vec![
6449            Some(1.0_f32),
6450            Some(1.0_f32),
6451            None,
6452            None,
6453            Some(2.5_f32),
6454        ]));
6455        ree_write_read_roundtrip(ree, flat);
6456    }
6457
6458    #[test]
6459    fn ree_sliced() {
6460        // A sliced (non-zero offset) REE array: verify that get_physical_index
6461        // correctly accounts for the logical offset when expanding.
6462        // Full array: run_ends [3, 5, 7] → [a,a,a, b,b, c,c]
6463        // After slice(2, 5) the logical view is [a, b, b, c, c].
6464        let full: ArrayRef = Arc::new(
6465            RunArray::try_new(
6466                &Int32Array::from(vec![3, 5, 7]),
6467                &StringArray::from(vec!["a", "b", "c"]),
6468            )
6469            .unwrap(),
6470        );
6471        let sliced = full.slice(2, 5);
6472        let flat: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "b", "c", "c"]));
6473        ree_write_read_roundtrip(sliced, flat);
6474    }
6475
6476    #[test]
6477    #[cfg_attr(miri, ignore)] // Takes too long
6478    fn test_number_distinct_values_exact_count() {
6479        // 50 distinct Int32 values repeated across 100k rows, with every 7th row null.
6480        // Nulls must not be counted as a distinct value.
6481        let cardinality = 50u32;
6482        let array: ArrayRef = Arc::new(Int32Array::from_iter((0..100_000u32).map(|i| {
6483            if i % 7 == 0 {
6484                None
6485            } else {
6486                Some((i % cardinality) as i32)
6487            }
6488        })));
6489        let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, true)]));
6490        let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
6491
6492        let props = WriterProperties::builder()
6493            .set_write_row_group_number_distinct_values(true)
6494            .build();
6495        let mut buf = Vec::new();
6496        let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), Some(props)).unwrap();
6497        writer.write(&batch).unwrap();
6498        let metadata = writer.close().unwrap();
6499
6500        let count = metadata
6501            .row_group(0)
6502            .column(0)
6503            .statistics()
6504            .and_then(|s| s.distinct_count_opt())
6505            .expect("distinct_count should be set");
6506        // Must equal cardinality exactly; nulls must not inflate the count.
6507        assert_eq!(count, cardinality as u64);
6508    }
6509
6510    #[test]
6511    fn test_number_distinct_values_view_types() {
6512        // 5 distinct values repeated across 30 rows, with every 4th row null.
6513        // Verifies Utf8View is counted correctly (BinaryView shares the same code path).
6514        let cardinality = 5u32;
6515        let distinct_strings = ["alpha", "beta", "gamma", "delta", "epsilon"];
6516
6517        let string_view_col: ArrayRef = Arc::new(StringViewArray::from_iter((0..30u32).map(|i| {
6518            if i % 4 == 0 {
6519                None
6520            } else {
6521                Some(distinct_strings[(i % cardinality) as usize])
6522            }
6523        })));
6524
6525        let schema = Arc::new(Schema::new(vec![Field::new(
6526            "string_view_col",
6527            DataType::Utf8View,
6528            true,
6529        )]));
6530        let batch = RecordBatch::try_new(schema, vec![string_view_col]).unwrap();
6531
6532        let props = WriterProperties::builder()
6533            .set_write_row_group_number_distinct_values(true)
6534            .build();
6535        let mut parquet_bytes = Vec::new();
6536        let mut writer =
6537            ArrowWriter::try_new(&mut parquet_bytes, batch.schema(), Some(props)).unwrap();
6538        writer.write(&batch).unwrap();
6539        let metadata = writer.close().unwrap();
6540
6541        let distinct_count = metadata
6542            .row_group(0)
6543            .column(0)
6544            .statistics()
6545            .and_then(|s| s.distinct_count_opt())
6546            .expect("distinct_count should be set for Utf8View column");
6547        assert_eq!(distinct_count, cardinality as u64);
6548    }
6549
6550    #[test]
6551    fn test_number_distinct_values_not_written_by_default() {
6552        let array: ArrayRef = Arc::new(Int32Array::from_iter_values(0..100));
6553        let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, false)]));
6554        let batch = RecordBatch::try_new(schema, vec![array]).unwrap();
6555
6556        let mut buf = Vec::new();
6557        let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), None).unwrap();
6558        writer.write(&batch).unwrap();
6559        let metadata = writer.close().unwrap();
6560
6561        let count = metadata
6562            .row_group(0)
6563            .column(0)
6564            .statistics()
6565            .and_then(|s| s.distinct_count_opt());
6566        assert!(count.is_none());
6567    }
6568
6569    #[test]
6570    fn ree_struct_with_ree_child() {
6571        // Struct with a REE string field and a REE int field — confirms
6572        // recursion visits every child and each collapses to the right leaf type.
6573        let run_ends = Int32Array::from(vec![2i32, 3, 5]);
6574
6575        let col_a: ArrayRef = Arc::new(
6576            RunArray::try_new(
6577                &run_ends,
6578                &StringArray::from(vec![Some("foo"), None, Some("bar")]),
6579            )
6580            .unwrap(),
6581        );
6582        let col_b: ArrayRef = Arc::new(
6583            RunArray::try_new(&run_ends, &Int32Array::from(vec![Some(1), None, Some(2)])).unwrap(),
6584        );
6585
6586        let struct_array: ArrayRef = Arc::new(StructArray::new(
6587            Fields::from(vec![
6588                Field::new("a", col_a.data_type().clone(), true),
6589                Field::new("b", col_b.data_type().clone(), true),
6590            ]),
6591            vec![col_a, col_b],
6592            None,
6593        ));
6594
6595        let schema = Arc::new(Schema::new(vec![Field::new(
6596            "row",
6597            struct_array.data_type().clone(),
6598            true,
6599        )]));
6600        let batch = RecordBatch::try_new(schema.clone(), vec![struct_array]).unwrap();
6601
6602        let mut buf = Vec::new();
6603        let mut writer = ArrowWriter::try_new(&mut buf, schema, None).unwrap();
6604        writer.write(&batch).unwrap();
6605        let metadata = writer.close().unwrap();
6606
6607        let parquet_schema = metadata.file_metadata().schema_descr();
6608        assert_eq!(parquet_schema.num_columns(), 2);
6609        assert_eq!(
6610            parquet_schema.column(0).physical_type(),
6611            crate::basic::Type::BYTE_ARRAY
6612        );
6613        assert_eq!(parquet_schema.column(0).path().string(), "row.a");
6614        assert_eq!(
6615            parquet_schema.column(1).physical_type(),
6616            crate::basic::Type::INT32
6617        );
6618        assert_eq!(parquet_schema.column(1).path().string(), "row.b");
6619    }
6620}