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